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