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 {
1692 return Ok(());
1693 };
1694 let message = format!(
1695 "Background run {status}.\n- run_id: {}\n- title: {}\n- tool: {}\n- summary: {}\n- result_path: {}\n- log_path: {}",
1696 self.run_id,
1697 self.display_name,
1698 self.tool_name,
1699 summary,
1700 self.result_path,
1701 self.log_path
1702 );
1703 platform_store
1704 .send_message(self.context.session_id, &message)
1705 .await
1706 }
1707
1708 async fn write_text_file(&self, path: &str, content: &str) -> Result<()> {
1709 let file_store = self.context.file_store.as_ref().ok_or_else(|| {
1710 anyhow::anyhow!(
1711 "background run {} cannot persist artifact {} because no session file store is configured",
1712 self.run_id,
1713 path
1714 )
1715 })?;
1716
1717 ensure_directory(file_store.as_ref(), self.context.session_id, "/.background").await?;
1718 let run_dir = format!("/.background/{}", self.run_id);
1719 ensure_directory(file_store.as_ref(), self.context.session_id, &run_dir).await?;
1720 file_store
1721 .write_file(self.context.session_id, path, content, "text")
1722 .await?;
1723 Ok(())
1724 }
1725}
1726
1727#[async_trait]
1728impl BackgroundEventSink for SessionBackgroundSink {
1729 async fn status(&self, message: &str) -> Result<()> {
1730 let mut state = self.state.lock().await;
1731 state.status_text = message.to_string();
1732 drop(state);
1733 self.mirror_task(crate::session_task::SessionTaskUpdate {
1734 state_detail: Some(message.to_string()),
1735 ..Default::default()
1736 })
1737 .await;
1738 Ok(())
1739 }
1740
1741 async fn output(&self, stream: &str, delta: &str) -> Result<()> {
1742 let mut state = self.state.lock().await;
1743 if !delta.is_empty() {
1744 let prefix = format!("[{stream}] ");
1745 state.output_tail.push_str(&prefix);
1746 state.output_tail.push_str(delta);
1747 Self::append_to_output_log(&mut state, &prefix, delta);
1748 if state.output_tail.chars().count() > 2048 {
1749 state.output_tail = state
1750 .output_tail
1751 .chars()
1752 .rev()
1753 .take(2048)
1754 .collect::<Vec<_>>()
1755 .into_iter()
1756 .rev()
1757 .collect();
1758 }
1759 }
1760 Ok(())
1761 }
1762
1763 async fn progress(&self, progress: BackgroundProgress) -> Result<()> {
1764 let mut state = self.state.lock().await;
1765 state.progress = Some(progress.clone());
1766 drop(state);
1767 self.mirror_task(crate::session_task::SessionTaskUpdate {
1768 progress: Some(progress),
1769 ..Default::default()
1770 })
1771 .await;
1772 Ok(())
1773 }
1774}
1775
1776impl SessionBackgroundSink {
1777 fn append_to_output_log(state: &mut SessionBackgroundState, prefix: &str, delta: &str) {
1778 if state.output_log_chars >= MAX_BACKGROUND_OUTPUT_LOG_CHARS {
1779 state.output_log_truncated = true;
1780 return;
1781 }
1782
1783 let chunk = format!("{prefix}{delta}");
1784 let remaining = MAX_BACKGROUND_OUTPUT_LOG_CHARS - state.output_log_chars;
1785 let chunk_chars = chunk.chars().count();
1786
1787 if chunk_chars <= remaining {
1788 state.output_log.push_str(&chunk);
1789 state.output_log_chars += chunk_chars;
1790 return;
1791 }
1792
1793 let truncated_chunk: String = chunk.chars().take(remaining).collect();
1794 state.output_log.push_str(&truncated_chunk);
1795 state.output_log_chars += truncated_chunk.chars().count();
1796 state.output_log_truncated = true;
1797 }
1798
1799 fn final_output_log(state: &SessionBackgroundState) -> String {
1800 if !state.output_log_truncated {
1801 return state.output_log.clone();
1802 }
1803
1804 format!(
1805 "{}\n[system] background output truncated at {} characters\n",
1806 state.output_log, MAX_BACKGROUND_OUTPUT_LOG_CHARS
1807 )
1808 }
1809}
1810
1811const BACKGROUND_CANCEL_SENTINEL: &str = "__everruns_background_cancel__";
1815
1816fn is_canceled_outcome(
1820 outcome: &std::result::Result<BackgroundOutcome, ToolExecutionResult>,
1821) -> bool {
1822 matches!(outcome, Err(ToolExecutionResult::ToolError(msg)) if msg == BACKGROUND_CANCEL_SENTINEL)
1823}
1824
1825pub(crate) async fn reattach_background_run(
1841 task: &crate::session_task::SessionTask,
1842 context: &crate::traits::ToolContext,
1843) -> crate::error::Result<()> {
1844 if context.file_store.is_none() {
1847 return Err(crate::error::AgentLoopError::tool(
1848 "file store not available; cannot re-attach background run",
1849 ));
1850 }
1851 if context.session_task_registry.is_none() {
1852 return Err(crate::error::AgentLoopError::tool(
1853 "task registry not available; cannot re-attach background run",
1854 ));
1855 }
1856
1857 let tool_name: String = task
1858 .spec
1859 .get("tool")
1860 .and_then(|v| v.as_str())
1861 .filter(|s| !s.is_empty())
1862 .map(str::to_owned)
1863 .ok_or_else(|| {
1864 crate::error::AgentLoopError::tool(
1865 "background_tool spec missing 'tool' field; cannot re-attach",
1866 )
1867 })?;
1868
1869 let tool_args = task
1870 .spec
1871 .get("arguments")
1872 .cloned()
1873 .unwrap_or_else(|| serde_json::Value::Object(Default::default()));
1874
1875 let registry = std::sync::Arc::new(ToolRegistry::with_defaults());
1876
1877 let Some(tool) = registry.get(&tool_name).cloned() else {
1878 return Err(crate::error::AgentLoopError::tool(format!(
1879 "tool '{tool_name}' not found in built-in registry; cannot re-attach"
1880 )));
1881 };
1882
1883 if tool.as_background_executable().is_none() {
1884 return Err(crate::error::AgentLoopError::tool(format!(
1885 "tool '{tool_name}' does not support background execution; cannot re-attach"
1886 )));
1887 }
1888
1889 let hints = tool.hints();
1892 if !hints.idempotent.unwrap_or(false) && !hints.readonly.unwrap_or(false) {
1893 return Err(crate::error::AgentLoopError::tool(format!(
1894 "tool '{tool_name}' is not idempotent or readonly; re-attach declined",
1895 )));
1896 }
1897
1898 let background_run_permit = try_acquire_background_run_permit(task.session_id)
1901 .map_err(crate::error::AgentLoopError::tool)?;
1902
1903 let signal_on_completion = task
1906 .spec
1907 .get("signal_on_completion")
1908 .and_then(|v| v.as_bool())
1909 .unwrap_or(true);
1910
1911 let run_id = format!("bg_{}", uuid::Uuid::now_v7().simple());
1912 let artifact_dir = format!("/.background/{run_id}");
1913 let log_path = format!("{artifact_dir}/output.log");
1914 let result_path = format!("{artifact_dir}/result.json");
1915
1916 let task_id = task.id.clone();
1917 let task_attempt = task.attempt;
1918 let session_id = task.session_id;
1919
1920 let sink_context = context.clone().with_tool_registry(registry);
1921 let sink = std::sync::Arc::new(SessionBackgroundSink::new(
1922 sink_context.clone(),
1923 run_id.clone(),
1924 task.display_name.clone(),
1925 tool_name.to_string(),
1926 log_path,
1927 result_path,
1928 signal_on_completion,
1929 Some(task_id.clone()),
1930 ));
1931
1932 let cancel_registry = context.session_task_registry.clone();
1933 let run_id_for_log = run_id.clone();
1934
1935 tokio::spawn(async move {
1936 let _background_run_permit = background_run_permit;
1938 let _ = sink.status("Re-attaching").await;
1939
1940 let outcome: std::result::Result<BackgroundOutcome, ToolExecutionResult> =
1941 match (cancel_registry.as_ref(), Some(task_id.as_str())) {
1942 (Some(registry), Some(task_id_str)) => {
1943 let registry = registry.clone();
1944 let task_id_str = task_id_str.to_string();
1945 let tool_fut = async {
1946 match tool.as_background_executable() {
1947 Some(bg) => {
1948 bg.execute_background(tool_args, sink_context.clone(), sink.clone())
1949 .await
1950 }
1951 None => Err(ToolExecutionResult::tool_error(format!(
1952 "tool '{tool_name}' lost background support during re-attach"
1953 ))),
1954 }
1955 };
1956 let watch_fut = async {
1957 loop {
1958 tokio::time::sleep(std::time::Duration::from_secs(2)).await;
1959 let _ = registry
1960 .update(
1961 session_id,
1962 &task_id_str,
1963 crate::session_task::SessionTaskUpdate {
1964 heartbeat_at: Some(chrono::Utc::now()),
1965 expected_attempt: Some(task_attempt),
1966 ..Default::default()
1967 },
1968 )
1969 .await;
1970 if let Ok(Some(t)) = registry.get(session_id, &task_id_str).await
1971 && t.cancel_requested_at.is_some()
1972 {
1973 break;
1974 }
1975 }
1976 };
1977 tokio::select! {
1978 result = tool_fut => result,
1979 () = watch_fut => Err(ToolExecutionResult::ToolError(
1980 BACKGROUND_CANCEL_SENTINEL.to_string(),
1981 )),
1982 }
1983 }
1984 _ => match tool.as_background_executable() {
1985 Some(bg) => {
1986 bg.execute_background(tool_args, sink_context, sink.clone())
1987 .await
1988 }
1989 None => Err(ToolExecutionResult::tool_error(format!(
1990 "tool '{tool_name}' lost background support during re-attach"
1991 ))),
1992 },
1993 };
1994
1995 let finalize_result = if is_canceled_outcome(&outcome) {
1996 sink.finalize_canceled().await
1997 } else {
1998 sink.finalize(outcome).await
1999 };
2000 if let Err(err) = finalize_result {
2001 tracing::warn!(
2002 run_id = run_id_for_log,
2003 error = %err,
2004 "Background run re-attach finalization failed"
2005 );
2006 }
2007 });
2008
2009 Ok(())
2010}
2011
2012async fn ensure_directory(
2013 file_store: &dyn crate::traits::SessionFileSystem,
2014 session_id: crate::SessionId,
2015 path: &str,
2016) -> Result<()> {
2017 if let Some(entry) = file_store.stat_file(session_id, path).await? {
2018 if entry.is_directory {
2019 return Ok(());
2020 }
2021 return Err(anyhow::anyhow!("path exists but is not a directory: {path}").into());
2022 }
2023 let _ = file_store.create_directory(session_id, path).await?;
2024 Ok(())
2025}
2026
2027pub struct FailingTool {
2029 error_message: String,
2030 use_internal_error: bool,
2031}
2032
2033impl FailingTool {
2034 pub fn with_tool_error(message: impl Into<String>) -> Self {
2036 Self {
2037 error_message: message.into(),
2038 use_internal_error: false,
2039 }
2040 }
2041
2042 pub fn with_internal_error(message: impl Into<String>) -> Self {
2044 Self {
2045 error_message: message.into(),
2046 use_internal_error: true,
2047 }
2048 }
2049}
2050
2051impl Default for FailingTool {
2052 fn default() -> Self {
2053 Self::with_tool_error("Tool execution failed")
2054 }
2055}
2056
2057#[async_trait]
2058impl Tool for FailingTool {
2059 fn name(&self) -> &str {
2060 "failing_tool"
2061 }
2062
2063 fn display_name(&self) -> Option<&str> {
2064 Some("Failing Tool")
2065 }
2066
2067 fn description(&self) -> &str {
2068 "A tool that always fails (for testing error handling)"
2069 }
2070
2071 fn parameters_schema(&self) -> Value {
2072 serde_json::json!({
2073 "type": "object",
2074 "properties": {},
2075 "additionalProperties": false
2076 })
2077 }
2078
2079 fn hints(&self) -> ToolHints {
2080 ToolHints::default()
2081 .with_readonly(true)
2082 .with_idempotent(true)
2083 }
2084
2085 async fn execute(&self, _arguments: Value) -> ToolExecutionResult {
2086 if self.use_internal_error {
2087 ToolExecutionResult::internal_error_msg(&self.error_message)
2088 } else {
2089 ToolExecutionResult::tool_error(&self.error_message)
2090 }
2091 }
2092}
2093
2094#[cfg(test)]
2099mod tests {
2100 use super::*;
2101 use crate::capabilities::GetCurrentTimeTool;
2102 use crate::platform_store::PlatformStore;
2103 use crate::session_file::{FileInfo, FileStat, SessionFile};
2104 use crate::session_task::SessionTaskRegistry;
2105 use crate::traits::{SessionFileSystem, SessionScheduleStore};
2106 use crate::typed_id::{HarnessId, SessionId};
2107 use crate::{AgentId, KeyInfo, PlatformMessage, SecretInfo};
2108 use async_trait::async_trait;
2109 use std::sync::{
2110 Arc as StdArc, Mutex,
2111 atomic::{AtomicBool, Ordering},
2112 };
2113
2114 #[derive(Default)]
2115 struct TestBackgroundTool;
2116
2117 #[async_trait]
2118 impl BackgroundExecutableTool for TestBackgroundTool {
2119 async fn execute_background(
2120 &self,
2121 arguments: Value,
2122 _context: ToolContext,
2123 sink: Arc<dyn BackgroundEventSink>,
2124 ) -> std::result::Result<BackgroundOutcome, ToolExecutionResult> {
2125 sink.status("Waiting for test result")
2126 .await
2127 .map_err(ToolExecutionResult::internal_error)?;
2128 sink.output("stdout", "hello from background")
2129 .await
2130 .map_err(ToolExecutionResult::internal_error)?;
2131 sink.progress(BackgroundProgress {
2132 current: Some(1),
2133 total: Some(1),
2134 unit: Some("step".to_string()),
2135 label: Some("done".to_string()),
2136 })
2137 .await
2138 .map_err(ToolExecutionResult::internal_error)?;
2139
2140 Ok(BackgroundOutcome {
2141 summary: arguments["summary"].as_str().unwrap_or("done").to_string(),
2142 result: json!({"ok": true}),
2143 raw_output: None,
2144 })
2145 }
2146 }
2147
2148 #[async_trait]
2149 impl Tool for TestBackgroundTool {
2150 fn name(&self) -> &str {
2151 "test_background"
2152 }
2153
2154 fn display_name(&self) -> Option<&str> {
2155 Some("Test Background")
2156 }
2157
2158 fn description(&self) -> &str {
2159 "test tool"
2160 }
2161
2162 fn parameters_schema(&self) -> Value {
2163 json!({
2164 "type": "object",
2165 "properties": {
2166 "summary": { "type": "string" }
2167 }
2168 })
2169 }
2170
2171 async fn execute(&self, _arguments: Value) -> ToolExecutionResult {
2172 ToolExecutionResult::tool_error("foreground unsupported")
2173 }
2174
2175 fn hints(&self) -> ToolHints {
2176 ToolHints::default().with_supports_background(true)
2177 }
2178
2179 fn as_background_executable(&self) -> Option<&dyn BackgroundExecutableTool> {
2180 Some(self)
2181 }
2182 }
2183
2184 #[derive(Default)]
2185 struct TestFailingBackgroundTool;
2186
2187 #[async_trait]
2188 impl BackgroundExecutableTool for TestFailingBackgroundTool {
2189 async fn execute_background(
2190 &self,
2191 _arguments: Value,
2192 _context: ToolContext,
2193 sink: Arc<dyn BackgroundEventSink>,
2194 ) -> std::result::Result<BackgroundOutcome, ToolExecutionResult> {
2195 sink.status("Running failing test")
2196 .await
2197 .map_err(ToolExecutionResult::internal_error)?;
2198 sink.output("stderr", "background failed")
2199 .await
2200 .map_err(ToolExecutionResult::internal_error)?;
2201 Err(ToolExecutionResult::tool_error("boom"))
2202 }
2203 }
2204
2205 #[async_trait]
2206 impl Tool for TestFailingBackgroundTool {
2207 fn name(&self) -> &str {
2208 "test_background_fail"
2209 }
2210
2211 fn display_name(&self) -> Option<&str> {
2212 Some("Test Background Fail")
2213 }
2214
2215 fn description(&self) -> &str {
2216 "failing background test tool"
2217 }
2218
2219 fn parameters_schema(&self) -> Value {
2220 json!({
2221 "type": "object",
2222 "properties": {}
2223 })
2224 }
2225
2226 async fn execute(&self, _arguments: Value) -> ToolExecutionResult {
2227 ToolExecutionResult::tool_error("foreground unsupported")
2228 }
2229
2230 fn hints(&self) -> ToolHints {
2231 ToolHints::default().with_supports_background(true)
2232 }
2233
2234 fn as_background_executable(&self) -> Option<&dyn BackgroundExecutableTool> {
2235 Some(self)
2236 }
2237 }
2238
2239 #[derive(Default)]
2240 struct TestLargeOutputBackgroundTool;
2241
2242 #[async_trait]
2243 impl BackgroundExecutableTool for TestLargeOutputBackgroundTool {
2244 async fn execute_background(
2245 &self,
2246 _arguments: Value,
2247 _context: ToolContext,
2248 sink: Arc<dyn BackgroundEventSink>,
2249 ) -> std::result::Result<BackgroundOutcome, ToolExecutionResult> {
2250 let large_chunk = "x".repeat(MAX_BACKGROUND_OUTPUT_LOG_CHARS + 4096);
2251 sink.output("stdout", &large_chunk)
2252 .await
2253 .map_err(ToolExecutionResult::internal_error)?;
2254 Ok(BackgroundOutcome {
2255 summary: "large output complete".to_string(),
2256 result: json!({"ok": true}),
2257 raw_output: None,
2258 })
2259 }
2260 }
2261
2262 #[async_trait]
2263 impl Tool for TestLargeOutputBackgroundTool {
2264 fn name(&self) -> &str {
2265 "test_background_large_output"
2266 }
2267
2268 fn display_name(&self) -> Option<&str> {
2269 Some("Test Background Large Output")
2270 }
2271
2272 fn description(&self) -> &str {
2273 "background test tool with huge output"
2274 }
2275
2276 fn parameters_schema(&self) -> Value {
2277 json!({
2278 "type": "object",
2279 "properties": {}
2280 })
2281 }
2282
2283 async fn execute(&self, _arguments: Value) -> ToolExecutionResult {
2284 ToolExecutionResult::tool_error("foreground unsupported")
2285 }
2286
2287 fn hints(&self) -> ToolHints {
2288 ToolHints::default().with_supports_background(true)
2289 }
2290
2291 fn as_background_executable(&self) -> Option<&dyn BackgroundExecutableTool> {
2292 Some(self)
2293 }
2294 }
2295
2296 struct BlockingBackgroundTool {
2297 release: StdArc<AtomicBool>,
2298 }
2299
2300 #[async_trait]
2301 impl BackgroundExecutableTool for BlockingBackgroundTool {
2302 async fn execute_background(
2303 &self,
2304 _arguments: Value,
2305 _context: ToolContext,
2306 sink: Arc<dyn BackgroundEventSink>,
2307 ) -> std::result::Result<BackgroundOutcome, ToolExecutionResult> {
2308 sink.status("Blocking until released")
2309 .await
2310 .map_err(ToolExecutionResult::internal_error)?;
2311 while !self.release.load(Ordering::SeqCst) {
2312 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
2313 }
2314 Ok(BackgroundOutcome {
2315 summary: "released".to_string(),
2316 result: json!({"ok": true}),
2317 raw_output: None,
2318 })
2319 }
2320 }
2321
2322 #[async_trait]
2323 impl Tool for BlockingBackgroundTool {
2324 fn name(&self) -> &str {
2325 "test_background_blocking"
2326 }
2327
2328 fn display_name(&self) -> Option<&str> {
2329 Some("Test Background Blocking")
2330 }
2331
2332 fn description(&self) -> &str {
2333 "background test tool that waits for test release"
2334 }
2335
2336 fn parameters_schema(&self) -> Value {
2337 json!({
2338 "type": "object",
2339 "properties": {}
2340 })
2341 }
2342
2343 async fn execute(&self, _arguments: Value) -> ToolExecutionResult {
2344 ToolExecutionResult::tool_error("foreground unsupported")
2345 }
2346
2347 fn hints(&self) -> ToolHints {
2348 ToolHints::default().with_supports_background(true)
2349 }
2350
2351 fn as_background_executable(&self) -> Option<&dyn BackgroundExecutableTool> {
2352 Some(self)
2353 }
2354 }
2355
2356 #[derive(Default)]
2357 struct TestFileStore {
2358 files: Mutex<HashMap<String, SessionFile>>,
2359 }
2360
2361 #[async_trait]
2362 impl crate::traits::SessionFileSystem for TestFileStore {
2363 fn is_mount_resolver(&self) -> bool {
2364 false
2365 }
2366
2367 async fn read_file(
2368 &self,
2369 _session_id: SessionId,
2370 path: &str,
2371 ) -> crate::Result<Option<SessionFile>> {
2372 Ok(self.files.lock().unwrap().get(path).cloned())
2373 }
2374
2375 async fn write_file(
2376 &self,
2377 session_id: SessionId,
2378 path: &str,
2379 content: &str,
2380 encoding: &str,
2381 ) -> crate::Result<SessionFile> {
2382 let now = chrono::Utc::now();
2383 let file = SessionFile {
2384 id: uuid::Uuid::now_v7(),
2385 session_id: session_id.uuid(),
2386 path: path.to_string(),
2387 name: FileInfo::name_from_path(path),
2388 content: Some(content.to_string()),
2389 encoding: encoding.to_string(),
2390 is_directory: false,
2391 is_readonly: false,
2392 size_bytes: content.len() as i64,
2393 created_at: now,
2394 updated_at: now,
2395 };
2396 self.files
2397 .lock()
2398 .unwrap()
2399 .insert(path.to_string(), file.clone());
2400 Ok(file)
2401 }
2402
2403 async fn delete_file(
2404 &self,
2405 _session_id: SessionId,
2406 _path: &str,
2407 _recursive: bool,
2408 ) -> crate::Result<bool> {
2409 Ok(false)
2410 }
2411
2412 async fn list_directory(
2413 &self,
2414 _session_id: SessionId,
2415 _path: &str,
2416 ) -> crate::Result<Vec<FileInfo>> {
2417 Ok(Vec::new())
2418 }
2419
2420 async fn stat_file(
2421 &self,
2422 _session_id: SessionId,
2423 path: &str,
2424 ) -> crate::Result<Option<FileStat>> {
2425 let file = self.files.lock().unwrap().get(path).cloned();
2426 Ok(file.map(|entry| FileStat {
2427 path: entry.path,
2428 name: entry.name,
2429 is_directory: entry.is_directory,
2430 is_readonly: entry.is_readonly,
2431 size_bytes: entry.size_bytes,
2432 created_at: entry.created_at,
2433 updated_at: entry.updated_at,
2434 }))
2435 }
2436
2437 async fn grep_files(
2438 &self,
2439 _session_id: SessionId,
2440 _pattern: &str,
2441 _path_pattern: Option<&str>,
2442 ) -> crate::Result<Vec<crate::session_file::GrepMatch>> {
2443 Ok(Vec::new())
2444 }
2445
2446 async fn create_directory(
2447 &self,
2448 session_id: SessionId,
2449 path: &str,
2450 ) -> crate::Result<FileInfo> {
2451 let now = chrono::Utc::now();
2452 let id = uuid::Uuid::now_v7();
2453 let dir = SessionFile {
2454 id,
2455 session_id: session_id.uuid(),
2456 path: path.to_string(),
2457 name: FileInfo::name_from_path(path),
2458 content: None,
2459 encoding: "text".to_string(),
2460 is_directory: true,
2461 is_readonly: false,
2462 size_bytes: 0,
2463 created_at: now,
2464 updated_at: now,
2465 };
2466 self.files.lock().unwrap().insert(path.to_string(), dir);
2467 Ok(FileInfo {
2468 id,
2469 session_id: session_id.uuid(),
2470 path: path.to_string(),
2471 name: FileInfo::name_from_path(path),
2472 is_directory: true,
2473 is_readonly: false,
2474 size_bytes: 0,
2475 created_at: now,
2476 updated_at: now,
2477 })
2478 }
2479 }
2480
2481 #[derive(Default)]
2482 struct TestPlatformStore {
2483 sent_messages: Mutex<Vec<String>>,
2484 }
2485
2486 #[async_trait]
2487 impl PlatformStore for TestPlatformStore {
2488 async fn list_harnesses(&self) -> crate::Result<Vec<crate::Harness>> {
2489 Ok(Vec::new())
2490 }
2491 async fn get_harness(&self, _id: HarnessId) -> crate::Result<Option<crate::Harness>> {
2492 Ok(None)
2493 }
2494 async fn create_harness(
2495 &self,
2496 _name: &str,
2497 _display_name: Option<&str>,
2498 _description: Option<&str>,
2499 _system_prompt: Option<&str>,
2500 _parent_harness_id: Option<HarnessId>,
2501 _capabilities: &[String],
2502 ) -> crate::Result<crate::Harness> {
2503 unreachable!()
2504 }
2505 async fn update_harness(
2506 &self,
2507 _id: HarnessId,
2508 _name: Option<&str>,
2509 _display_name: Option<&str>,
2510 _description: Option<&str>,
2511 _system_prompt: Option<&str>,
2512 _parent_harness_id: Option<Option<HarnessId>>,
2513 ) -> crate::Result<crate::Harness> {
2514 unreachable!()
2515 }
2516 async fn delete_harness(&self, _id: HarnessId) -> crate::Result<()> {
2517 Ok(())
2518 }
2519 async fn copy_harness(
2520 &self,
2521 _id: HarnessId,
2522 _new_name: Option<&str>,
2523 ) -> crate::Result<crate::Harness> {
2524 unreachable!()
2525 }
2526 async fn list_agents(&self) -> crate::Result<Vec<crate::Agent>> {
2527 Ok(Vec::new())
2528 }
2529 async fn get_agent_by_id(&self, _id: AgentId) -> crate::Result<Option<crate::Agent>> {
2530 Ok(None)
2531 }
2532 async fn create_agent(
2533 &self,
2534 _name: &str,
2535 _display_name: Option<&str>,
2536 _description: Option<&str>,
2537 _system_prompt: &str,
2538 _capabilities: &[String],
2539 ) -> crate::Result<crate::Agent> {
2540 unreachable!()
2541 }
2542 async fn update_agent(
2543 &self,
2544 _id: AgentId,
2545 _name: Option<&str>,
2546 _display_name: Option<&str>,
2547 _description: Option<&str>,
2548 _system_prompt: Option<&str>,
2549 ) -> crate::Result<crate::Agent> {
2550 unreachable!()
2551 }
2552 async fn delete_agent(&self, _id: AgentId) -> crate::Result<()> {
2553 Ok(())
2554 }
2555 async fn list_apps(
2556 &self,
2557 _search: Option<&str>,
2558 _include_archived: bool,
2559 ) -> crate::Result<Vec<crate::App>> {
2560 Ok(Vec::new())
2561 }
2562 async fn get_app(&self, _id: crate::AppId) -> crate::Result<Option<crate::App>> {
2563 Ok(None)
2564 }
2565 async fn create_app(
2566 &self,
2567 _name: &str,
2568 _description: Option<&str>,
2569 _harness_id: HarnessId,
2570 _agent_id: Option<AgentId>,
2571 _agent_identity_id: Option<crate::AgentIdentityId>,
2572 _channel_type: Option<crate::ChannelType>,
2573 _channel_config: Option<&serde_json::Value>,
2574 ) -> crate::Result<crate::App> {
2575 unreachable!()
2576 }
2577 async fn update_app(
2578 &self,
2579 _id: crate::AppId,
2580 _name: Option<&str>,
2581 _description: Option<&str>,
2582 _harness_id: Option<HarnessId>,
2583 _agent_id: Option<AgentId>,
2584 _agent_identity_id: Option<Option<crate::AgentIdentityId>>,
2585 ) -> crate::Result<crate::App> {
2586 unreachable!()
2587 }
2588 async fn delete_app(&self, _id: crate::AppId) -> crate::Result<()> {
2589 Ok(())
2590 }
2591 async fn destroy_app(&self, _id: crate::AppId) -> crate::Result<()> {
2592 Ok(())
2593 }
2594 async fn publish_app(&self, _id: crate::AppId) -> crate::Result<crate::App> {
2595 unreachable!()
2596 }
2597 async fn unpublish_app(&self, _id: crate::AppId) -> crate::Result<crate::App> {
2598 unreachable!()
2599 }
2600 async fn add_app_channel(
2601 &self,
2602 _app_id: crate::AppId,
2603 _channel_type: crate::ChannelType,
2604 _channel_config: Option<&serde_json::Value>,
2605 _enabled: Option<bool>,
2606 ) -> crate::Result<crate::AppChannel> {
2607 unreachable!()
2608 }
2609 async fn update_app_channel(
2610 &self,
2611 _app_id: crate::AppId,
2612 _channel_id: crate::AppChannelId,
2613 _channel_type: Option<crate::ChannelType>,
2614 _channel_config: Option<&serde_json::Value>,
2615 _enabled: Option<bool>,
2616 ) -> crate::Result<crate::AppChannel> {
2617 unreachable!()
2618 }
2619 async fn delete_app_channel(
2620 &self,
2621 _app_id: crate::AppId,
2622 _channel_id: crate::AppChannelId,
2623 ) -> crate::Result<()> {
2624 Ok(())
2625 }
2626 async fn list_sessions(
2627 &self,
2628 _limit: Option<usize>,
2629 _agent_id: Option<AgentId>,
2630 ) -> crate::Result<Vec<crate::Session>> {
2631 Ok(Vec::new())
2632 }
2633 async fn create_session(
2634 &self,
2635 _harness_id: HarnessId,
2636 _agent_id: Option<AgentId>,
2637 _title: Option<&str>,
2638 _locale: Option<&str>,
2639 _blueprint_id: Option<&str>,
2640 _blueprint_config: Option<&serde_json::Value>,
2641 _parent_session_id: Option<SessionId>,
2642 ) -> crate::Result<crate::Session> {
2643 unreachable!()
2644 }
2645 async fn get_session_by_id(&self, _id: SessionId) -> crate::Result<Option<crate::Session>> {
2646 Ok(None)
2647 }
2648 async fn add_agent_session_participant(
2649 &self,
2650 _session_id: SessionId,
2651 _agent_id: AgentId,
2652 ) -> crate::Result<crate::SessionParticipant> {
2653 unreachable!()
2654 }
2655 async fn get_session_context_report(
2656 &self,
2657 id: SessionId,
2658 ) -> crate::Result<crate::SessionContextReport> {
2659 Ok(crate::SessionContextReport {
2660 session_id: id.to_string(),
2661 model: "llmsim".to_string(),
2662 context_window_tokens: None,
2663 estimated_input_tokens: 0,
2664 sections: vec![],
2665 contributions: vec![],
2666 cumulative_usage: None,
2667 })
2668 }
2669 async fn delete_session(&self, _id: SessionId) -> crate::Result<()> {
2670 Ok(())
2671 }
2672 async fn send_message(&self, _session_id: SessionId, content: &str) -> crate::Result<()> {
2673 self.sent_messages.lock().unwrap().push(content.to_string());
2674 Ok(())
2675 }
2676 async fn get_messages(
2677 &self,
2678 _session_id: SessionId,
2679 _limit: Option<usize>,
2680 ) -> crate::Result<Vec<PlatformMessage>> {
2681 Ok(Vec::new())
2682 }
2683 async fn wait_for_idle(
2684 &self,
2685 _session_id: SessionId,
2686 _timeout_secs: Option<u64>,
2687 ) -> crate::Result<String> {
2688 Ok("idle".to_string())
2689 }
2690 async fn list_capabilities(
2691 &self,
2692 _search: Option<&str>,
2693 ) -> crate::Result<Vec<crate::CapabilityInfo>> {
2694 Ok(Vec::new())
2695 }
2696 fn base_url(&self) -> &str {
2697 "http://localhost:9300"
2698 }
2699 }
2700
2701 #[derive(Default)]
2702 struct NoopStorageStore;
2703
2704 #[derive(Default)]
2705 struct TestScheduleStore {
2706 schedules: Mutex<Vec<crate::session_schedule::SessionSchedule>>,
2707 }
2708
2709 #[async_trait]
2710 impl crate::traits::SessionScheduleStore for TestScheduleStore {
2711 async fn create_schedule(
2712 &self,
2713 session_id: SessionId,
2714 description: String,
2715 cron_expression: Option<String>,
2716 scheduled_at: Option<chrono::DateTime<chrono::Utc>>,
2717 timezone: String,
2718 ) -> crate::Result<crate::session_schedule::SessionSchedule> {
2719 let schedule = crate::session_schedule::SessionSchedule {
2720 id: crate::typed_id::ScheduleId::new(),
2721 session_id,
2722 owner_principal_id: crate::PrincipalId::from_seed(1),
2723 resolved_owner_user_id: None,
2724 owner: None,
2725 effective_owner: None,
2726 description,
2727 cron_expression: cron_expression.clone(),
2728 scheduled_at,
2729 timezone,
2730 enabled: true,
2731 schedule_type: crate::session_schedule::SessionSchedule::derive_type(
2732 &cron_expression,
2733 ),
2734 next_trigger_at: Some(chrono::Utc::now() + chrono::Duration::minutes(10)),
2735 last_triggered_at: None,
2736 trigger_count: 0,
2737 created_at: chrono::Utc::now(),
2738 updated_at: chrono::Utc::now(),
2739 };
2740 self.schedules.lock().unwrap().push(schedule.clone());
2741 Ok(schedule)
2742 }
2743
2744 async fn cancel_schedule(
2745 &self,
2746 _session_id: SessionId,
2747 schedule_id: crate::ScheduleId,
2748 ) -> crate::Result<crate::session_schedule::SessionSchedule> {
2749 let mut schedules = self.schedules.lock().unwrap();
2750 let schedule = schedules
2751 .iter_mut()
2752 .find(|schedule| schedule.id == schedule_id)
2753 .ok_or_else(|| crate::AgentLoopError::tool("Schedule not found".to_string()))?;
2754 schedule.enabled = false;
2755 Ok(schedule.clone())
2756 }
2757
2758 async fn list_schedules(
2759 &self,
2760 session_id: SessionId,
2761 ) -> crate::Result<Vec<crate::session_schedule::SessionSchedule>> {
2762 Ok(self
2763 .schedules
2764 .lock()
2765 .unwrap()
2766 .iter()
2767 .filter(|schedule| schedule.session_id == session_id)
2768 .cloned()
2769 .collect())
2770 }
2771
2772 async fn count_active_schedules(&self, session_id: SessionId) -> crate::Result<u32> {
2773 Ok(self
2774 .schedules
2775 .lock()
2776 .unwrap()
2777 .iter()
2778 .filter(|schedule| schedule.session_id == session_id && schedule.enabled)
2779 .count() as u32)
2780 }
2781
2782 async fn count_active_org_schedules(&self) -> crate::Result<u32> {
2783 Ok(self
2785 .schedules
2786 .lock()
2787 .unwrap()
2788 .iter()
2789 .filter(|schedule| schedule.enabled)
2790 .count() as u32)
2791 }
2792 }
2793
2794 #[async_trait]
2795 impl crate::traits::SessionStorageStore for NoopStorageStore {
2796 async fn set_value(
2797 &self,
2798 _session_id: SessionId,
2799 _key: &str,
2800 _value: &str,
2801 ) -> crate::Result<()> {
2802 Ok(())
2803 }
2804 async fn get_value(
2805 &self,
2806 _session_id: SessionId,
2807 _key: &str,
2808 ) -> crate::Result<Option<String>> {
2809 Ok(None)
2810 }
2811 async fn delete_value(&self, _session_id: SessionId, _key: &str) -> crate::Result<bool> {
2812 Ok(false)
2813 }
2814 async fn list_keys(&self, _session_id: SessionId) -> crate::Result<Vec<KeyInfo>> {
2815 Ok(Vec::new())
2816 }
2817 async fn set_secret(
2818 &self,
2819 _session_id: SessionId,
2820 _name: &str,
2821 _value: &str,
2822 ) -> crate::Result<()> {
2823 Ok(())
2824 }
2825 async fn get_secret(
2826 &self,
2827 _session_id: SessionId,
2828 _name: &str,
2829 ) -> crate::Result<Option<String>> {
2830 Ok(None)
2831 }
2832 async fn delete_secret(&self, _session_id: SessionId, _name: &str) -> crate::Result<bool> {
2833 Ok(false)
2834 }
2835 async fn list_secrets(&self, _session_id: SessionId) -> crate::Result<Vec<SecretInfo>> {
2836 Ok(Vec::new())
2837 }
2838 }
2839
2840 #[tokio::test]
2841 async fn test_echo_tool() {
2842 let tool = EchoTool;
2843
2844 let result = tool
2845 .execute(serde_json::json!({"message": "Hello, world!"}))
2846 .await;
2847
2848 if let ToolExecutionResult::Success(value) = result {
2849 assert_eq!(
2850 value.get("echoed").unwrap().as_str().unwrap(),
2851 "Hello, world!"
2852 );
2853 assert_eq!(value.get("length").unwrap().as_u64().unwrap(), 13);
2854 } else {
2855 panic!("Expected success");
2856 }
2857 }
2858
2859 #[tokio::test]
2860 async fn test_failing_tool_with_tool_error() {
2861 let tool = FailingTool::with_tool_error("Something went wrong");
2862
2863 let result = tool.execute(serde_json::json!({})).await;
2864
2865 if let ToolExecutionResult::ToolError(msg) = result {
2866 assert_eq!(msg, "Something went wrong");
2867 } else {
2868 panic!("Expected tool error");
2869 }
2870 }
2871
2872 #[tokio::test]
2873 async fn test_failing_tool_with_internal_error() {
2874 let tool = FailingTool::with_internal_error("Database connection failed");
2875
2876 let result = tool.execute(serde_json::json!({})).await;
2877
2878 if let ToolExecutionResult::InternalError(err) = result {
2879 assert_eq!(err.message, "Database connection failed");
2880 } else {
2881 panic!("Expected internal error");
2882 }
2883 }
2884
2885 #[tokio::test]
2886 async fn test_tool_result_conversion() {
2887 let result = ToolExecutionResult::success(serde_json::json!({"value": 42}));
2889 let tool_result = result.into_tool_result("call_1", "test_tool");
2890 assert!(tool_result.error.is_none());
2891 assert_eq!(tool_result.result.unwrap()["value"], 42);
2892
2893 let result = ToolExecutionResult::tool_error("Invalid input");
2895 let tool_result = result.into_tool_result("call_2", "test_tool");
2896 assert_eq!(tool_result.error.as_deref(), Some("Invalid input"));
2897 assert_eq!(
2898 tool_result.result.unwrap(),
2899 serde_json::json!({"error": "Invalid input"})
2900 );
2901
2902 let result = ToolExecutionResult::internal_error_msg("Secret database error");
2904 let tool_result = result.into_tool_result("call_3", "test_tool");
2905 assert_eq!(
2906 tool_result.error.as_deref(),
2907 Some("An internal error occurred while executing the tool")
2908 );
2909 assert_eq!(
2910 tool_result.result.unwrap(),
2911 serde_json::json!({"error": "An internal error occurred while executing the tool"})
2912 );
2913 }
2914
2915 #[tokio::test]
2916 async fn test_tool_registry() {
2917 let mut registry = ToolRegistry::new();
2918 registry.register(GetCurrentTimeTool);
2919 registry.register(EchoTool);
2920
2921 assert_eq!(registry.len(), 2);
2922 assert!(registry.has("get_current_time"));
2923 assert!(registry.has("echo"));
2924 assert!(!registry.has("nonexistent"));
2925
2926 let definitions = registry.tool_definitions();
2927 assert_eq!(definitions.len(), 2);
2928 }
2929
2930 #[tokio::test]
2931 async fn test_tool_registry_builder() {
2932 let registry = ToolRegistry::builder()
2933 .tool(GetCurrentTimeTool)
2934 .tool(EchoTool)
2935 .build();
2936
2937 assert_eq!(registry.len(), 2);
2938 }
2939
2940 #[test]
2941 fn test_tool_display_name_in_definition() {
2942 let tool = GetCurrentTimeTool;
2944 assert_eq!(tool.display_name(), Some("Get Current Time"));
2945
2946 let def = tool.to_definition();
2947 assert_eq!(def.display_name(), Some("Get Current Time"));
2948 }
2949
2950 #[test]
2951 fn test_success_with_raw_output_object_preserves_shape() {
2952 let res = ToolExecutionResult::success_with_raw_output(
2953 serde_json::json!({"stdout": "hello"}),
2954 "raw stdout bytes".to_string(),
2955 );
2956 let tr = res.into_tool_result("call_1", "demo");
2957 assert_eq!(tr.result.as_ref().unwrap()["stdout"], "hello");
2958 assert!(
2959 tr.result
2960 .as_ref()
2961 .unwrap()
2962 .as_object()
2963 .unwrap()
2964 .get("_raw_output")
2965 .is_none(),
2966 "sidecar key must not leak to the LLM-visible result"
2967 );
2968 assert_eq!(tr.raw_output.as_deref(), Some("raw stdout bytes"));
2969 }
2970
2971 #[test]
2972 fn test_success_with_raw_output_scalar_unwraps_to_string() {
2973 let res = ToolExecutionResult::success_with_raw_output(
2974 "compact summary".to_string(),
2975 "full output bytes".to_string(),
2976 );
2977 let tr = res.into_tool_result("call_1", "demo");
2978 assert_eq!(
2979 tr.result,
2980 Some(serde_json::Value::String("compact summary".into()))
2981 );
2982 assert_eq!(tr.raw_output.as_deref(), Some("full output bytes"));
2983 }
2984
2985 #[test]
2986 fn test_success_result_with_raw_output_scalar_key_is_not_unwrapped() {
2987 let res = ToolExecutionResult::success(
2988 serde_json::json!({"_raw_output_scalar": "user_value", "kept": true}),
2989 );
2990 let tr = res.into_tool_result("call_1", "demo");
2991 assert_eq!(
2992 tr.result,
2993 Some(serde_json::json!({"_raw_output_scalar": "user_value", "kept": true}))
2994 );
2995 assert_eq!(tr.raw_output, None);
2996 }
2997
2998 #[test]
2999 fn test_success_result_with_only_raw_output_scalar_key_is_not_unwrapped() {
3000 let res = ToolExecutionResult::success(serde_json::json!({"_raw_output_scalar": "v"}));
3003 let tr = res.into_tool_result("call_1", "demo");
3004 assert_eq!(
3005 tr.result,
3006 Some(serde_json::json!({"_raw_output_scalar": "v"}))
3007 );
3008 assert_eq!(tr.raw_output, None);
3009 }
3010
3011 #[test]
3012 fn test_echo_tool_display_name() {
3013 let tool = EchoTool;
3014 assert_eq!(tool.display_name(), Some("Echo"));
3015
3016 let def = tool.to_definition();
3017 assert_eq!(def.display_name(), Some("Echo"));
3018 }
3019
3020 #[test]
3021 fn test_all_default_tools_have_display_names() {
3022 let registry = ToolRegistry::with_defaults();
3023 let definitions = registry.tool_definitions();
3024
3025 for def in &definitions {
3026 assert!(
3027 def.display_name().is_some(),
3028 "Tool '{}' should have a display_name",
3029 def.name()
3030 );
3031 }
3032 }
3033
3034 #[tokio::test]
3035 async fn test_tool_registry_as_executor() {
3036 let mut registry = ToolRegistry::new();
3037 registry.register(EchoTool);
3038
3039 let tool_call = ToolCall {
3040 id: "call_1".to_string(),
3041 name: "echo".to_string(),
3042 arguments: serde_json::json!({"message": "test"}),
3043 };
3044
3045 let tool_def = registry.get("echo").unwrap().to_definition();
3046 let result = registry.execute(&tool_call, &tool_def).await.unwrap();
3047
3048 assert!(result.error.is_none());
3049 assert_eq!(result.result.unwrap()["echoed"], "test");
3050 }
3051
3052 #[test]
3053 fn test_tool_to_definition() {
3054 let tool = GetCurrentTimeTool;
3055 let def = tool.to_definition();
3056
3057 let ToolDefinition::Builtin(builtin) = def else {
3058 panic!("expected Builtin variant");
3059 };
3060 assert_eq!(builtin.name, "get_current_time");
3061 assert_eq!(builtin.policy, ToolPolicy::Auto);
3062 }
3063
3064 #[test]
3065 fn test_with_defaults_has_expected_tools() {
3066 let registry = ToolRegistry::with_defaults();
3067
3068 assert!(
3070 registry.has("get_current_time"),
3071 "should have get_current_time"
3072 );
3073 assert!(registry.has("echo"), "should have echo");
3074 assert!(
3077 !registry.has("spawn_background"),
3078 "spawn_background must NOT be in defaults — it comes from the \
3079 background_execution capability"
3080 );
3081 assert!(
3082 registry.has("report_progress"),
3083 "should have report_progress"
3084 );
3085
3086 assert!(registry.has("add"), "should have add");
3088 assert!(registry.has("subtract"), "should have subtract");
3089 assert!(registry.has("multiply"), "should have multiply");
3090 assert!(registry.has("divide"), "should have divide");
3091
3092 assert!(registry.has("get_weather"), "should have get_weather");
3094 assert!(registry.has("get_forecast"), "should have get_forecast");
3095
3096 assert!(registry.has("write_todos"), "should have write_todos");
3098
3099 assert!(registry.has("read_file"), "should have read_file");
3101 assert!(registry.has("write_file"), "should have write_file");
3102 assert!(registry.has("edit_file"), "should have edit_file");
3103 assert!(registry.has("list_directory"), "should have list_directory");
3104 assert!(registry.has("grep_files"), "should have grep_files");
3105 assert!(registry.has("delete_file"), "should have delete_file");
3106 assert!(registry.has("stat_file"), "should have stat_file");
3107
3108 assert!(registry.has("web_fetch"), "should have web_fetch");
3110
3111 assert_eq!(registry.len(), 18, "should have 18 default tools");
3113 }
3114
3115 #[tokio::test]
3116 async fn test_with_defaults_tools_are_executable() {
3117 let registry = ToolRegistry::with_defaults();
3118
3119 let tool_call = ToolCall {
3121 id: "call_1".to_string(),
3122 name: "echo".to_string(),
3123 arguments: serde_json::json!({"message": "hello from defaults"}),
3124 };
3125
3126 let tool_def = registry.get("echo").unwrap().to_definition();
3127 let result = registry.execute(&tool_call, &tool_def).await.unwrap();
3128
3129 assert!(result.error.is_none());
3130 assert_eq!(result.result.unwrap()["echoed"], "hello from defaults");
3131 }
3132
3133 #[tokio::test]
3134 async fn test_with_defaults_math_tools() {
3135 let registry = ToolRegistry::with_defaults();
3136
3137 let tool_call = ToolCall {
3139 id: "call_add".to_string(),
3140 name: "add".to_string(),
3141 arguments: serde_json::json!({"a": 5, "b": 3}),
3142 };
3143
3144 let tool_def = registry.get("add").unwrap().to_definition();
3145 let result = registry.execute(&tool_call, &tool_def).await.unwrap();
3146
3147 assert!(result.error.is_none());
3148 assert_eq!(result.result.unwrap()["result"].as_f64().unwrap(), 8.0);
3150 }
3151
3152 #[test]
3156 fn test_with_defaults_excludes_capability_only_tools() {
3157 let registry = ToolRegistry::with_defaults();
3158
3159 assert!(
3161 !registry.has("bash"),
3162 "bash must not be in defaults — it comes from bashkit_shell capability"
3163 );
3164 assert!(
3166 !registry.has("kv_store"),
3167 "kv_store must not be in defaults — it comes from session_storage capability"
3168 );
3169 assert!(
3173 !registry.has("spawn_background"),
3174 "spawn_background must not be in defaults — it comes from the \
3175 background_execution capability (auto-activated by tool hints)"
3176 );
3177 }
3178
3179 #[tokio::test]
3180 async fn test_spawn_background_executes_and_signals_session() {
3181 let session_id = SessionId::new();
3182 let file_store = Arc::new(TestFileStore::default());
3183 let platform_store = Arc::new(TestPlatformStore::default());
3184 let storage_store = Arc::new(NoopStorageStore);
3185 let task_registry = Arc::new(InMemoryTaskRegistry::default());
3186 let tool_registry = ToolRegistry::builder()
3187 .tool(SpawnBackgroundTool)
3188 .tool(TestBackgroundTool)
3189 .build();
3190
3191 let context = ToolContext::with_stores(session_id, file_store.clone(), storage_store)
3192 .with_tool_registry(Arc::new(tool_registry))
3193 .with_platform_store(platform_store.clone())
3194 .with_session_task_registry(task_registry.clone());
3195
3196 let tool = SpawnBackgroundTool;
3197 let result = tool
3198 .execute_with_context(
3199 json!({
3200 "tool": "test_background",
3201 "args": { "summary": "Background complete" }
3202 }),
3203 &context,
3204 )
3205 .await;
3206
3207 let ToolExecutionResult::Success(value) = result else {
3208 panic!("spawn_background should succeed");
3209 };
3210 let run_id = value["run_id"].as_str().unwrap().to_string();
3211 let task_id = value["task_id"].as_str().unwrap().to_string();
3212
3213 tokio::time::timeout(std::time::Duration::from_secs(2), async {
3214 loop {
3215 if let Ok(Some(task)) = task_registry.get(session_id, &task_id).await
3216 && task.state == crate::session_task::SessionTaskState::Succeeded
3217 {
3218 break task;
3219 }
3220 tokio::time::sleep(std::time::Duration::from_millis(20)).await;
3221 }
3222 })
3223 .await
3224 .expect("background run should complete");
3225 let _ = run_id; let messages = platform_store.sent_messages.lock().unwrap().clone();
3228 assert_eq!(messages.len(), 1);
3229 assert!(messages[0].contains("Background run completed"));
3230
3231 let log_file = file_store
3232 .read_file(session_id, &format!("/.background/{run_id}/output.log"))
3233 .await
3234 .unwrap()
3235 .expect("log file");
3236 assert!(
3237 log_file
3238 .content
3239 .as_deref()
3240 .unwrap_or_default()
3241 .contains("hello from background")
3242 );
3243 }
3244
3245 #[tokio::test]
3246 async fn test_spawn_background_persists_failure_artifacts() {
3247 let session_id = SessionId::new();
3248 let file_store = Arc::new(TestFileStore::default());
3249 let storage_store = Arc::new(NoopStorageStore);
3250 let task_registry = Arc::new(InMemoryTaskRegistry::default());
3251 let tool_registry = ToolRegistry::builder()
3252 .tool(SpawnBackgroundTool)
3253 .tool(TestFailingBackgroundTool)
3254 .build();
3255
3256 let context = ToolContext::with_stores(session_id, file_store.clone(), storage_store)
3257 .with_tool_registry(Arc::new(tool_registry))
3258 .with_session_task_registry(task_registry.clone());
3259
3260 let result = SpawnBackgroundTool
3261 .execute_with_context(
3262 json!({
3263 "tool": "test_background_fail",
3264 "args": {}
3265 }),
3266 &context,
3267 )
3268 .await;
3269
3270 let ToolExecutionResult::Success(value) = result else {
3271 panic!("spawn_background should succeed");
3272 };
3273 let run_id = value["run_id"].as_str().unwrap().to_string();
3274 let task_id = value["task_id"].as_str().unwrap().to_string();
3275
3276 tokio::time::timeout(std::time::Duration::from_secs(2), async {
3277 loop {
3278 if let Ok(Some(task)) = task_registry.get(session_id, &task_id).await
3279 && task.state == crate::session_task::SessionTaskState::Failed
3280 {
3281 break task;
3282 }
3283 tokio::time::sleep(std::time::Duration::from_millis(20)).await;
3284 }
3285 })
3286 .await
3287 .expect("background run should fail");
3288 let _ = run_id;
3289
3290 let log_file = file_store
3291 .read_file(session_id, &format!("/.background/{run_id}/output.log"))
3292 .await
3293 .unwrap()
3294 .expect("log file");
3295 assert!(
3296 log_file
3297 .content
3298 .as_deref()
3299 .unwrap_or_default()
3300 .contains("background failed")
3301 );
3302
3303 let result_file = file_store
3304 .read_file(session_id, &format!("/.background/{run_id}/result.json"))
3305 .await
3306 .unwrap()
3307 .expect("result file");
3308 let result_json: Value =
3309 serde_json::from_str(result_file.content.as_deref().unwrap_or_default())
3310 .expect("valid json");
3311 assert_eq!(result_json["status"], "failed");
3312 assert_eq!(result_json["error"], "boom");
3313 }
3314
3315 #[tokio::test]
3316 async fn test_spawn_background_rejects_when_session_active_run_limit_reached() {
3317 let session_id = SessionId::new();
3318 let file_store = Arc::new(TestFileStore::default());
3319 let storage_store = Arc::new(NoopStorageStore);
3320 let task_registry = Arc::new(InMemoryTaskRegistry::default());
3321 let release = StdArc::new(AtomicBool::new(false));
3322 let tool_registry = ToolRegistry::builder()
3323 .tool(SpawnBackgroundTool)
3324 .tool(BlockingBackgroundTool {
3325 release: release.clone(),
3326 })
3327 .build();
3328
3329 let context = ToolContext::with_stores(session_id, file_store, storage_store)
3330 .with_tool_registry(Arc::new(tool_registry))
3331 .with_session_task_registry(task_registry.clone());
3332
3333 let mut task_ids = Vec::new();
3334 for _ in 0..MAX_ACTIVE_BACKGROUND_RUNS_PER_SESSION {
3335 let result = SpawnBackgroundTool
3336 .execute_with_context(
3337 json!({
3338 "tool": "test_background_blocking",
3339 "args": {}
3340 }),
3341 &context,
3342 )
3343 .await;
3344
3345 let ToolExecutionResult::Success(value) = result else {
3346 panic!("background run below the session limit should start");
3347 };
3348 task_ids.push(value["task_id"].as_str().unwrap().to_string());
3349 }
3350
3351 tokio::time::timeout(std::time::Duration::from_secs(2), async {
3353 loop {
3354 let running = task_registry
3355 .list(
3356 session_id,
3357 Some(&crate::session_task::SessionTaskFilter {
3358 kind: Some(crate::session_task::TASK_KIND_BACKGROUND_TOOL.to_string()),
3359 state: Some(crate::session_task::SessionTaskState::Running),
3360 }),
3361 )
3362 .await
3363 .unwrap();
3364 if running.len() == MAX_ACTIVE_BACKGROUND_RUNS_PER_SESSION {
3365 break;
3366 }
3367 tokio::time::sleep(std::time::Duration::from_millis(20)).await;
3368 }
3369 })
3370 .await
3371 .expect("background runs should become running");
3372
3373 let result = SpawnBackgroundTool
3374 .execute_with_context(
3375 json!({
3376 "tool": "test_background_blocking",
3377 "args": {}
3378 }),
3379 &context,
3380 )
3381 .await;
3382
3383 let ToolExecutionResult::ToolError(message) = result else {
3384 release.store(true, Ordering::SeqCst);
3385 panic!("spawn_background should reject once the session limit is reached");
3386 };
3387 assert!(message.contains("active background runs per session"));
3388
3389 release.store(true, Ordering::SeqCst);
3390 tokio::time::timeout(std::time::Duration::from_secs(2), async {
3391 for task_id in task_ids {
3392 loop {
3393 if let Ok(Some(task)) = task_registry.get(session_id, &task_id).await
3394 && task.state == crate::session_task::SessionTaskState::Succeeded
3395 {
3396 break;
3397 }
3398 tokio::time::sleep(std::time::Duration::from_millis(20)).await;
3399 }
3400 }
3401 })
3402 .await
3403 .expect("blocking background runs should complete after release");
3404
3405 tokio::time::timeout(std::time::Duration::from_secs(1), async {
3410 loop {
3411 if !has_session_background_permits(session_id) {
3412 break;
3413 }
3414 tokio::time::sleep(std::time::Duration::from_millis(5)).await;
3415 }
3416 })
3417 .await
3418 .expect("completed background runs should prune their per-session permit cache entry");
3419 }
3420
3421 #[tokio::test]
3422 async fn test_spawn_background_requires_task_registry() {
3423 let session_id = SessionId::new();
3424 let file_store = Arc::new(TestFileStore::default());
3425 let storage_store = Arc::new(NoopStorageStore);
3426 let tool_registry = ToolRegistry::builder()
3427 .tool(SpawnBackgroundTool)
3428 .tool(TestBackgroundTool)
3429 .build();
3430
3431 let context = ToolContext::with_stores(session_id, file_store, storage_store)
3433 .with_tool_registry(Arc::new(tool_registry));
3434
3435 let result = SpawnBackgroundTool
3436 .execute_with_context(
3437 json!({
3438 "tool": "test_background",
3439 "args": {}
3440 }),
3441 &context,
3442 )
3443 .await;
3444
3445 let ToolExecutionResult::ToolError(message) = result else {
3446 panic!("spawn_background should reject missing task registry");
3447 };
3448 assert!(message.contains("Session task registry not available"));
3449 }
3450
3451 #[tokio::test]
3452 async fn test_spawn_background_requires_file_store() {
3453 let session_id = SessionId::new();
3454 let storage_store = Arc::new(NoopStorageStore);
3455 let task_registry = Arc::new(InMemoryTaskRegistry::default());
3456 let tool_registry = ToolRegistry::builder()
3457 .tool(SpawnBackgroundTool)
3458 .tool(TestBackgroundTool)
3459 .build();
3460
3461 let context = ToolContext::with_storage_store(session_id, storage_store)
3463 .with_tool_registry(Arc::new(tool_registry))
3464 .with_session_task_registry(task_registry);
3465
3466 let result = SpawnBackgroundTool
3467 .execute_with_context(
3468 json!({
3469 "tool": "test_background",
3470 "args": {}
3471 }),
3472 &context,
3473 )
3474 .await;
3475
3476 let ToolExecutionResult::ToolError(message) = result else {
3477 panic!("spawn_background should reject missing file store");
3478 };
3479 assert!(message.contains("Session file store not available"));
3480 }
3481
3482 #[tokio::test]
3483 async fn test_spawn_background_caps_output_log_size() {
3484 let session_id = SessionId::new();
3485 let file_store = Arc::new(TestFileStore::default());
3486 let storage_store = Arc::new(NoopStorageStore);
3487 let task_registry = Arc::new(InMemoryTaskRegistry::default());
3488 let tool_registry = ToolRegistry::builder()
3489 .tool(SpawnBackgroundTool)
3490 .tool(TestLargeOutputBackgroundTool)
3491 .build();
3492
3493 let context = ToolContext::with_stores(session_id, file_store.clone(), storage_store)
3494 .with_tool_registry(Arc::new(tool_registry))
3495 .with_session_task_registry(task_registry.clone());
3496
3497 let result = SpawnBackgroundTool
3498 .execute_with_context(
3499 json!({
3500 "tool": "test_background_large_output",
3501 "args": {}
3502 }),
3503 &context,
3504 )
3505 .await;
3506
3507 let ToolExecutionResult::Success(value) = result else {
3508 panic!("spawn_background should succeed");
3509 };
3510 let run_id = value["run_id"].as_str().unwrap().to_string();
3511 let task_id = value["task_id"].as_str().unwrap().to_string();
3512
3513 tokio::time::timeout(std::time::Duration::from_secs(2), async {
3514 loop {
3515 if let Ok(Some(task)) = task_registry.get(session_id, &task_id).await
3516 && task.state == crate::session_task::SessionTaskState::Succeeded
3517 {
3518 break;
3519 }
3520 tokio::time::sleep(std::time::Duration::from_millis(20)).await;
3521 }
3522 })
3523 .await
3524 .expect("background run should complete");
3525 let _ = run_id;
3526
3527 let log_content = file_store
3528 .read_file(session_id, &format!("/.background/{run_id}/output.log"))
3529 .await
3530 .unwrap()
3531 .expect("log file")
3532 .content
3533 .unwrap_or_default();
3534
3535 assert!(log_content.contains("[system] background output truncated"));
3536 assert!(log_content.chars().count() <= MAX_BACKGROUND_OUTPUT_LOG_CHARS + 128);
3537 }
3538
3539 #[tokio::test]
3540 async fn test_spawn_background_can_create_scheduled_monitor() {
3541 let session_id = SessionId::new();
3542 let schedule_store = Arc::new(TestScheduleStore::default());
3543 let storage_store = Arc::new(NoopStorageStore);
3544 let tool_registry = ToolRegistry::builder()
3545 .tool(SpawnBackgroundTool)
3546 .tool(TestBackgroundTool)
3547 .build();
3548
3549 let context = ToolContext::with_storage_store(session_id, storage_store)
3550 .with_tool_registry(Arc::new(tool_registry))
3551 .with_schedule_store(schedule_store.clone());
3552
3553 let result = SpawnBackgroundTool
3554 .execute_with_context(
3555 json!({
3556 "tool": "test_background",
3557 "title": "Watch PR 1319",
3558 "args": { "summary": "Background complete" },
3559 "schedule": {
3560 "cron_expression": "*/10 * * * *",
3561 "timezone": "America/Chicago"
3562 }
3563 }),
3564 &context,
3565 )
3566 .await;
3567
3568 let ToolExecutionResult::Success(value) = result else {
3569 panic!("spawn_background should create a schedule: {result:?}");
3570 };
3571
3572 assert_eq!(value["status"], "scheduled");
3573 assert_eq!(value["title"], "Watch PR 1319");
3574 assert_eq!(value["cron_expression"], "*/10 * * * *");
3575 assert_eq!(value["timezone"], "America/Chicago");
3576
3577 let schedules = schedule_store.list_schedules(session_id).await.unwrap();
3578 assert_eq!(schedules.len(), 1);
3579 assert_eq!(
3580 schedules[0].cron_expression.as_deref(),
3581 Some("*/10 * * * *")
3582 );
3583 assert!(schedules[0].description.contains("Monitor: Watch PR 1319"));
3584 assert!(
3585 schedules[0]
3586 .description
3587 .contains("\"summary\": \"Background complete\"")
3588 );
3589 }
3590
3591 #[tokio::test]
3592 async fn test_spawn_background_rejects_invalid_scheduled_at() {
3593 let session_id = SessionId::new();
3594 let storage_store = Arc::new(NoopStorageStore);
3595 let tool_registry = ToolRegistry::builder()
3596 .tool(SpawnBackgroundTool)
3597 .tool(TestBackgroundTool)
3598 .build();
3599 let context = ToolContext::with_storage_store(session_id, storage_store)
3600 .with_tool_registry(Arc::new(tool_registry));
3601
3602 let result = SpawnBackgroundTool
3603 .execute_with_context(
3604 json!({
3605 "tool": "test_background",
3606 "args": {},
3607 "schedule": {
3608 "scheduled_at": "tomorrow at noon"
3609 }
3610 }),
3611 &context,
3612 )
3613 .await;
3614
3615 let ToolExecutionResult::ToolError(message) = result else {
3616 panic!("spawn_background should reject invalid scheduled_at");
3617 };
3618 assert!(message.contains("scheduled_at must be RFC3339"));
3619 }
3620
3621 #[tokio::test]
3622 async fn test_spawn_background_rejects_ambiguous_schedule_shape() {
3623 let session_id = SessionId::new();
3624 let storage_store = Arc::new(NoopStorageStore);
3625 let tool_registry = ToolRegistry::builder()
3626 .tool(SpawnBackgroundTool)
3627 .tool(TestBackgroundTool)
3628 .build();
3629 let context = ToolContext::with_storage_store(session_id, storage_store)
3630 .with_tool_registry(Arc::new(tool_registry));
3631
3632 let result = SpawnBackgroundTool
3633 .execute_with_context(
3634 json!({
3635 "tool": "test_background",
3636 "args": {},
3637 "schedule": {
3638 "cron_expression": "*/10 * * * *",
3639 "scheduled_at": "2026-04-16T15:30:00Z"
3640 }
3641 }),
3642 &context,
3643 )
3644 .await;
3645
3646 let ToolExecutionResult::ToolError(message) = result else {
3647 panic!("spawn_background should reject ambiguous schedule shape");
3648 };
3649 assert!(message.contains("must not include both cron_expression and scheduled_at"));
3650 }
3651
3652 #[test]
3657 fn test_is_canceled_outcome_detects_sentinel() {
3658 let sentinel: std::result::Result<BackgroundOutcome, ToolExecutionResult> = Err(
3660 ToolExecutionResult::ToolError(BACKGROUND_CANCEL_SENTINEL.to_string()),
3661 );
3662 assert!(is_canceled_outcome(&sentinel));
3663 }
3664
3665 #[test]
3666 fn test_is_canceled_outcome_does_not_match_other_errors() {
3667 let other_err: std::result::Result<BackgroundOutcome, ToolExecutionResult> =
3668 Err(ToolExecutionResult::ToolError("boom".to_string()));
3669 assert!(!is_canceled_outcome(&other_err));
3670
3671 let success: std::result::Result<BackgroundOutcome, ToolExecutionResult> =
3672 Ok(BackgroundOutcome {
3673 summary: "ok".to_string(),
3674 result: json!({"ok": true}),
3675 raw_output: None,
3676 });
3677 assert!(!is_canceled_outcome(&success));
3678 }
3679
3680 #[derive(Default)]
3683 struct SleepingBackgroundTool;
3684
3685 #[async_trait]
3686 impl BackgroundExecutableTool for SleepingBackgroundTool {
3687 async fn execute_background(
3688 &self,
3689 _arguments: Value,
3690 _context: ToolContext,
3691 sink: Arc<dyn BackgroundEventSink>,
3692 ) -> std::result::Result<BackgroundOutcome, ToolExecutionResult> {
3693 sink.status("Sleeping forever")
3694 .await
3695 .map_err(ToolExecutionResult::internal_error)?;
3696 tokio::time::sleep(std::time::Duration::from_secs(3600)).await;
3698 Ok(BackgroundOutcome {
3699 summary: "should not reach here".to_string(),
3700 result: json!({}),
3701 raw_output: None,
3702 })
3703 }
3704 }
3705
3706 #[async_trait]
3707 impl Tool for SleepingBackgroundTool {
3708 fn name(&self) -> &str {
3709 "test_background_sleeping"
3710 }
3711
3712 fn display_name(&self) -> Option<&str> {
3713 Some("Test Background Sleeping")
3714 }
3715
3716 fn description(&self) -> &str {
3717 "background test tool that sleeps indefinitely"
3718 }
3719
3720 fn parameters_schema(&self) -> Value {
3721 json!({
3722 "type": "object",
3723 "properties": {}
3724 })
3725 }
3726
3727 async fn execute(&self, _arguments: Value) -> ToolExecutionResult {
3728 ToolExecutionResult::tool_error("foreground unsupported")
3729 }
3730
3731 fn hints(&self) -> ToolHints {
3732 ToolHints::default().with_supports_background(true)
3733 }
3734
3735 fn as_background_executable(&self) -> Option<&dyn BackgroundExecutableTool> {
3736 Some(self)
3737 }
3738 }
3739
3740 #[derive(Default)]
3744 struct InMemoryTaskRegistry {
3745 tasks: Mutex<HashMap<String, crate::session_task::SessionTask>>,
3746 }
3747
3748 #[async_trait]
3749 impl crate::session_task::SessionTaskRegistry for InMemoryTaskRegistry {
3750 async fn create(
3751 &self,
3752 input: crate::session_task::CreateSessionTask,
3753 ) -> crate::Result<crate::session_task::SessionTask> {
3754 let mut tasks = self.tasks.lock().unwrap();
3755 if let Some(id) = &input.id
3756 && let Some(existing) = tasks.get(id)
3757 {
3758 return Ok(existing.clone());
3759 }
3760 let task = crate::session_task::new_session_task(input, chrono::Utc::now());
3761 tasks.insert(task.id.clone(), task.clone());
3762 Ok(task)
3763 }
3764
3765 async fn update(
3766 &self,
3767 _session_id: SessionId,
3768 task_id: &str,
3769 update: crate::session_task::SessionTaskUpdate,
3770 ) -> crate::Result<Option<crate::session_task::SessionTask>> {
3771 let mut tasks = self.tasks.lock().unwrap();
3772 let Some(task) = tasks.get_mut(task_id) else {
3773 return Ok(None);
3774 };
3775 crate::session_task::apply_task_update(task, update, chrono::Utc::now());
3776 Ok(Some(task.clone()))
3777 }
3778
3779 async fn get(
3780 &self,
3781 _session_id: SessionId,
3782 task_id: &str,
3783 ) -> crate::Result<Option<crate::session_task::SessionTask>> {
3784 Ok(self.tasks.lock().unwrap().get(task_id).cloned())
3785 }
3786
3787 async fn list(
3788 &self,
3789 _session_id: SessionId,
3790 filter: Option<&crate::session_task::SessionTaskFilter>,
3791 ) -> crate::Result<Vec<crate::session_task::SessionTask>> {
3792 let tasks = self.tasks.lock().unwrap();
3793 Ok(tasks
3794 .values()
3795 .filter(|task| {
3796 filter.is_none_or(|f| {
3797 f.kind.as_deref().is_none_or(|kind| task.kind == kind)
3798 && f.state.is_none_or(|state| task.state == state)
3799 })
3800 })
3801 .cloned()
3802 .collect())
3803 }
3804
3805 async fn request_cancel(
3806 &self,
3807 _session_id: SessionId,
3808 task_id: &str,
3809 ) -> crate::Result<Option<crate::session_task::SessionTask>> {
3810 let mut tasks = self.tasks.lock().unwrap();
3811 let Some(task) = tasks.get_mut(task_id) else {
3812 return Ok(None);
3813 };
3814 task.cancel_requested_at
3815 .get_or_insert_with(chrono::Utc::now);
3816 task.updated_at = chrono::Utc::now();
3817 Ok(Some(task.clone()))
3818 }
3819
3820 async fn record_message(
3821 &self,
3822 _session_id: SessionId,
3823 task_id: &str,
3824 message: crate::session_task::NewTaskMessage,
3825 ) -> crate::Result<crate::session_task::TaskMessage> {
3826 let tasks = self.tasks.lock().unwrap();
3827 let _task = tasks
3828 .get(task_id)
3829 .ok_or_else(|| crate::AgentLoopError::tool(format!("no task {task_id}")))?;
3830 Ok(crate::session_task::TaskMessage {
3831 id: crate::session_task::generate_task_message_id(),
3832 task_id: task_id.to_string(),
3833 direction: message.direction,
3834 content: message.content,
3835 in_reply_to: message.in_reply_to,
3836 created_at: chrono::Utc::now(),
3837 })
3838 }
3839
3840 async fn list_messages(
3841 &self,
3842 _session_id: SessionId,
3843 _task_id: &str,
3844 _limit: Option<u32>,
3845 _after_id: Option<&str>,
3846 ) -> crate::Result<Vec<crate::session_task::TaskMessage>> {
3847 Ok(Vec::new())
3848 }
3849 }
3850
3851 #[tokio::test]
3854 async fn test_cancel_background_run_via_task_registry() {
3855 let session_id = SessionId::new();
3856 let file_store = Arc::new(TestFileStore::default());
3857 let storage_store = Arc::new(NoopStorageStore);
3858 let task_registry = Arc::new(InMemoryTaskRegistry::default());
3859
3860 let tool_registry = ToolRegistry::builder()
3861 .tool(SpawnBackgroundTool)
3862 .tool(SleepingBackgroundTool)
3863 .build();
3864
3865 let context = ToolContext::with_stores(session_id, file_store.clone(), storage_store)
3866 .with_tool_registry(Arc::new(tool_registry))
3867 .with_session_task_registry(task_registry.clone());
3868
3869 let result = SpawnBackgroundTool
3870 .execute_with_context(
3871 json!({
3872 "tool": "test_background_sleeping",
3873 "args": {},
3874 "signal_on_completion": false
3875 }),
3876 &context,
3877 )
3878 .await;
3879
3880 let ToolExecutionResult::Success(value) = result else {
3881 panic!("spawn_background should succeed");
3882 };
3883 let run_id = value["run_id"].as_str().unwrap().to_string();
3884 let task_id = value["task_id"].as_str().unwrap().to_string();
3885
3886 tokio::time::timeout(std::time::Duration::from_secs(5), async {
3888 loop {
3889 if let Ok(Some(task)) = task_registry.get(session_id, &task_id).await
3891 && task.heartbeat_at.is_some()
3892 {
3893 break;
3894 }
3895 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
3896 }
3897 })
3898 .await
3899 .expect("background run should start and send at least one heartbeat");
3900
3901 task_registry
3903 .request_cancel(session_id, &task_id)
3904 .await
3905 .expect("request_cancel should succeed");
3906
3907 tokio::time::timeout(std::time::Duration::from_secs(10), async {
3909 loop {
3910 if let Ok(Some(task)) = task_registry.get(session_id, &task_id).await
3911 && task.state == crate::session_task::SessionTaskState::Canceled
3912 {
3913 break task;
3914 }
3915 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
3916 }
3917 })
3918 .await
3919 .expect("background task should reach Canceled state");
3920
3921 let result_file = file_store
3923 .read_file(session_id, &format!("/.background/{run_id}/result.json"))
3924 .await
3925 .unwrap()
3926 .expect("result.json should exist");
3927 let result_json: Value =
3928 serde_json::from_str(result_file.content.as_deref().unwrap_or_default())
3929 .expect("valid json");
3930 assert_eq!(result_json["status"], "canceled");
3931
3932 let log_file = file_store
3933 .read_file(session_id, &format!("/.background/{run_id}/output.log"))
3934 .await
3935 .unwrap()
3936 .expect("output.log should exist");
3937 assert!(
3938 log_file
3939 .content
3940 .as_deref()
3941 .unwrap_or_default()
3942 .contains("Canceled by request.")
3943 );
3944 }
3945
3946 fn make_reattach_task(spec: serde_json::Value) -> crate::session_task::SessionTask {
3951 use crate::session_task::{SessionTaskState, TaskLinks, TaskWakePolicy};
3952 crate::session_task::SessionTask {
3953 id: "t-reattach".to_string(),
3954 session_id: SessionId::new(),
3955 root_session_id: None,
3956 kind: crate::session_task::TASK_KIND_BACKGROUND_TOOL.to_string(),
3957 display_name: "Reattach test".to_string(),
3958 spec,
3959 state: SessionTaskState::Running,
3960 state_detail: None,
3961 progress: None,
3962 input_request: None,
3963 cancel_requested_at: None,
3964 summary: None,
3965 result_path: None,
3966 artifacts: vec![],
3967 error: None,
3968 attempt: 2,
3969 worker_id: None,
3970 heartbeat_at: None,
3971 links: TaskLinks::default(),
3972 wake_policy: TaskWakePolicy::Silent,
3973 created_at: chrono::Utc::now(),
3974 started_at: None,
3975 finished_at: None,
3976 updated_at: chrono::Utc::now(),
3977 }
3978 }
3979
3980 #[tokio::test]
3981 async fn reattach_fails_with_missing_file_store() {
3982 let session_id = SessionId::new();
3983 let task_registry = Arc::new(InMemoryTaskRegistry::default());
3985 let context =
3986 crate::traits::ToolContext::new(session_id).with_session_task_registry(task_registry);
3987 let task = make_reattach_task(serde_json::json!({
3988 "tool": "get_current_time",
3989 "arguments": {},
3990 "reattachable": true,
3991 "signal_on_completion": true,
3992 }));
3993 let err = reattach_background_run(&task, &context)
3994 .await
3995 .expect_err("should fail without file store");
3996 assert!(
3997 err.to_string().contains("file store"),
3998 "error should mention file store, got: {err}"
3999 );
4000 }
4001
4002 #[tokio::test]
4003 async fn reattach_fails_with_missing_task_registry() {
4004 let session_id = SessionId::new();
4005 let file_store = Arc::new(TestFileStore::default());
4006 let storage_store = Arc::new(NoopStorageStore);
4007 let context =
4009 crate::traits::ToolContext::with_stores(session_id, file_store, storage_store);
4010 let task = make_reattach_task(serde_json::json!({
4011 "tool": "get_current_time",
4012 "arguments": {},
4013 "reattachable": true,
4014 "signal_on_completion": true,
4015 }));
4016 let err = reattach_background_run(&task, &context)
4017 .await
4018 .expect_err("should fail without task registry");
4019 assert!(
4020 err.to_string().contains("task registry"),
4021 "error should mention task registry, got: {err}"
4022 );
4023 }
4024
4025 #[tokio::test]
4026 async fn reattach_fails_with_unknown_tool_name() {
4027 let session_id = SessionId::new();
4028 let file_store = Arc::new(TestFileStore::default());
4029 let storage_store = Arc::new(NoopStorageStore);
4030 let task_registry = Arc::new(InMemoryTaskRegistry::default());
4031 let context =
4032 crate::traits::ToolContext::with_stores(session_id, file_store, storage_store)
4033 .with_session_task_registry(task_registry);
4034 let task = make_reattach_task(serde_json::json!({
4036 "tool": "test_background",
4037 "arguments": {},
4038 "reattachable": true,
4039 "signal_on_completion": true,
4040 }));
4041 let err = reattach_background_run(&task, &context)
4042 .await
4043 .expect_err("should fail for unknown tool");
4044 assert!(
4045 err.to_string().contains("not found in built-in registry"),
4046 "error should mention built-in registry, got: {err}"
4047 );
4048 }
4049
4050 #[tokio::test]
4051 async fn reattach_fails_with_missing_tool_spec_field() {
4052 let session_id = SessionId::new();
4053 let file_store = Arc::new(TestFileStore::default());
4054 let storage_store = Arc::new(NoopStorageStore);
4055 let task_registry = Arc::new(InMemoryTaskRegistry::default());
4056 let context =
4057 crate::traits::ToolContext::with_stores(session_id, file_store, storage_store)
4058 .with_session_task_registry(task_registry);
4059 let task = make_reattach_task(serde_json::json!({ "reattachable": true }));
4061 let err = reattach_background_run(&task, &context)
4062 .await
4063 .expect_err("should fail with missing tool field");
4064 assert!(
4065 err.to_string().contains("missing 'tool' field"),
4066 "error should mention missing tool field, got: {err}"
4067 );
4068 }
4069}