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