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