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