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