Skip to main content

everruns_core/
tools.rs

1// Tool Abstraction for Agent Loop
2//
3// This module provides a high-level abstraction for tools that can be executed
4// by the agent loop. Tools are defined using the `Tool` trait and can be
5// registered with a `ToolRegistry` for use in the loop.
6//
7// Design decisions:
8// - Tools are defined via a trait for flexibility (function-style tools)
9// - ToolRegistry implements ToolExecutor for integration with the agent loop
10// - Error handling distinguishes between user-visible errors and internal errors
11// - Internal errors are logged but not exposed to the LLM (security)
12
13use 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
33/// Maximum active immediate background runs allowed for a single session.
34///
35/// This mirrors the scheduled monitor cap so model-visible background execution
36/// cannot be used to queue unbounded active worker jobs for one session.
37pub const MAX_ACTIVE_BACKGROUND_RUNS_PER_SESSION: usize = 5;
38
39/// Maximum active immediate background runs allowed in this worker process.
40///
41/// The per-session semaphore limits tenant/session abuse; this process-wide
42/// semaphore keeps concurrent sessions from exhausting worker-local resources.
43const 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
47/// Per-session semaphores that enforce `MAX_ACTIVE_BACKGROUND_RUNS_PER_SESSION`.
48///
49/// Using a semaphore rather than a DB count makes the check atomic: concurrent
50/// `spawn_background` calls for the same session cannot both slip through the
51/// guard (check-then-act race) because `try_acquire` is inherently atomic.
52static 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// ============================================================================
107// Tool Execution Result - Error Handling Contract
108// ============================================================================
109
110/// Image data returned by a tool alongside text results.
111///
112/// This allows tools (built-in or MCP) to return images that are sent
113/// to the LLM as native image content blocks, not stringified JSON.
114#[derive(Debug, Clone, Serialize, Deserialize)]
115pub struct ToolResultImage {
116    /// Base64-encoded image data
117    pub base64: String,
118    /// MIME type (e.g., "image/png", "image/jpeg")
119    pub media_type: String,
120}
121
122/// Result of a tool execution.
123///
124/// This enum distinguishes between different outcomes:
125/// - `Success`: Tool executed successfully, result is returned to LLM
126/// - `SuccessWithImages`: Successful execution with JSON result plus images
127/// - `ToolError`: Tool-level error that should be shown to the LLM
128///   (e.g., "City not found", "Invalid date format")
129/// - `InternalError`: System-level error that should NOT be exposed to the LLM
130///   (e.g., database connection failure, API key issues)
131///
132/// # Security
133///
134/// Internal errors are logged but replaced with a generic message when
135/// returned to the LLM. This prevents leaking sensitive information like
136/// database errors, API keys, or internal system details.
137#[derive(Debug)]
138pub enum ToolExecutionResult {
139    /// Successful execution with a JSON result
140    Success(Value),
141
142    /// Successful execution with a JSON result and images.
143    /// Images are sent to the LLM as native image content blocks
144    /// (not stringified JSON), enabling visual understanding.
145    SuccessWithImages {
146        result: Value,
147        images: Vec<ToolResultImage>,
148    },
149
150    /// Tool-level error that is safe to show to the LLM
151    ///
152    /// Use this for expected error conditions that the LLM should know about,
153    /// such as validation errors, resource not found, etc.
154    ToolError(String),
155
156    /// Internal/system error that should NOT be exposed to the LLM
157    ///
158    /// Use this for unexpected errors like network failures, database errors,
159    /// or other internal issues. The error details will be logged but replaced
160    /// with a generic message when returned to the LLM.
161    InternalError(ToolInternalError),
162
163    /// A user connection is required to execute this tool.
164    ///
165    /// Instead of returning an error, this signals that the workflow should
166    /// pause and ask the client to set up a connection for the given provider.
167    /// The UI renders an inline connection dialog; once the user saves (or
168    /// cancels), a tool result is submitted and execution resumes.
169    ConnectionRequired {
170        /// Connection provider id (e.g. "daytona", "brave_search")
171        provider: String,
172    },
173}
174
175impl ToolExecutionResult {
176    /// Create a successful result
177    pub fn success(value: impl Into<Value>) -> Self {
178        ToolExecutionResult::Success(value.into())
179    }
180
181    /// Create a successful result with pre-truncation raw output for VFS persistence.
182    /// The raw output is transferred to `ToolResult.raw_output` during `into_tool_result()`.
183    pub fn success_with_raw_output(value: impl Into<Value>, raw_output: String) -> Self {
184        let mut value = value.into();
185        // Embed raw output in a sidecar key — extracted in into_tool_result().
186        // Non-object values are wrapped in a scalar carrier so raw_output still
187        // flows through; the carrier is unwrapped on extraction.
188        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    /// Create a successful result with images
203    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    /// Create a tool-level error (safe to show to LLM)
211    pub fn tool_error(message: impl Into<String>) -> Self {
212        ToolExecutionResult::ToolError(message.into())
213    }
214
215    /// Create an internal error (will be hidden from LLM)
216    pub fn internal_error(error: impl std::error::Error + Send + Sync + 'static) -> Self {
217        ToolExecutionResult::InternalError(ToolInternalError::new(error))
218    }
219
220    /// Create an internal error from a string message
221    pub fn internal_error_msg(message: impl Into<String>) -> Self {
222        ToolExecutionResult::InternalError(ToolInternalError::from_message(message))
223    }
224
225    /// Signal that a user connection is required before this tool can execute.
226    pub fn connection_required(provider: impl Into<String>) -> Self {
227        ToolExecutionResult::ConnectionRequired {
228            provider: provider.into(),
229        }
230    }
231
232    /// Check if this is a successful result
233    pub fn is_success(&self) -> bool {
234        matches!(
235            self,
236            ToolExecutionResult::Success(_) | ToolExecutionResult::SuccessWithImages { .. }
237        )
238    }
239
240    /// Check if this is an error (either tool error or internal error)
241    pub fn is_error(&self) -> bool {
242        matches!(
243            self,
244            ToolExecutionResult::ToolError(_) | ToolExecutionResult::InternalError(_)
245        )
246    }
247
248    /// Check if this requires a user connection setup
249    pub fn is_connection_required(&self) -> bool {
250        matches!(self, ToolExecutionResult::ConnectionRequired { .. })
251    }
252
253    /// Convert to a ToolResult for the agent loop
254    ///
255    /// Both tool errors and internal errors are packaged as `{"error": "..."}` in the
256    /// result field. This provides a consistent contract where the result field always
257    /// contains the payload, and the agent loop continues the same way for all outcomes.
258    ///
259    /// Internal errors are logged but replaced with a generic message when returned.
260    pub fn into_tool_result(self, tool_call_id: &str, tool_name: &str) -> ToolResult {
261        match self {
262            ToolExecutionResult::Success(mut value) => {
263                // Extract sidecar raw output if present (from success_with_raw_output)
264                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                // Unwrap scalar carrier only when it matches the exact wrapper shape
269                // set by success_with_raw_output for non-object inputs.
270                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                // Log the full error details for debugging
313                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                // Return generic error message to LLM, packaged as {"error": "..."}
322                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/// Internal error details (logged but not exposed to LLM)
349#[derive(Debug)]
350pub struct ToolInternalError {
351    /// Error message for logging
352    pub message: String,
353    /// Optional source error
354    pub source: Option<Box<dyn std::error::Error + Send + Sync>>,
355}
356
357impl ToolInternalError {
358    /// Create from an error
359    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    /// Create from a string message
367    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// ============================================================================
403// Tool Trait - Core Tool Abstraction
404// ============================================================================
405
406/// Trait for implementing tools that can be executed by the agent loop.
407///
408/// # Example
409///
410/// ```ignore
411/// use async_trait::async_trait;
412/// use serde_json::{json, Value};
413///
414/// struct GetCurrentTime;
415///
416/// #[async_trait]
417/// impl Tool for GetCurrentTime {
418///     fn name(&self) -> &str {
419///         "get_current_time"
420///     }
421///
422///     fn description(&self) -> &str {
423///         "Get the current date and time"
424///     }
425///
426///     fn parameters_schema(&self) -> Value {
427///         json!({
428///             "type": "object",
429///             "properties": {
430///                 "timezone": {
431///                     "type": "string",
432///                     "description": "Timezone (e.g., 'UTC', 'America/New_York')"
433///                 }
434///             }
435///         })
436///     }
437///
438///     async fn execute(&self, arguments: Value) -> ToolExecutionResult {
439///         let timezone = arguments.get("timezone")
440///             .and_then(|v| v.as_str())
441///             .unwrap_or("UTC");
442///
443///         ToolExecutionResult::success(json!({
444///             "current_time": chrono::Utc::now().to_rfc3339(),
445///             "timezone": timezone
446///         }))
447///     }
448/// }
449/// ```
450#[async_trait]
451pub trait Tool: Send + Sync {
452    /// Returns the tool's unique name.
453    ///
454    /// This name is used by the LLM to invoke the tool and must be unique
455    /// within a ToolRegistry.
456    fn name(&self) -> &str;
457
458    /// Returns a human-readable display name for UI rendering.
459    ///
460    /// This name is shown to users in the UI instead of the technical tool name.
461    /// For example, "Get Current Time" instead of "get_current_time".
462    /// Returns None if no display name is set, in which case the UI may
463    /// fall back to the technical name.
464    fn display_name(&self) -> Option<&str> {
465        None
466    }
467
468    /// Returns a description of what the tool does.
469    ///
470    /// This description is provided to the LLM to help it understand
471    /// when and how to use the tool.
472    fn description(&self) -> &str;
473
474    /// Returns the JSON schema for the tool's parameters.
475    ///
476    /// This schema follows the JSON Schema specification and describes
477    /// the expected arguments for the tool. The LLM uses this to
478    /// generate valid tool calls.
479    fn parameters_schema(&self) -> Value;
480
481    /// Execute the tool with the given arguments.
482    ///
483    /// # Arguments
484    ///
485    /// * `arguments` - The arguments passed to the tool as a JSON value.
486    ///   These should conform to the schema returned by `parameters_schema()`.
487    ///
488    /// # Returns
489    ///
490    /// A `ToolExecutionResult` indicating success, tool error, or internal error.
491    async fn execute(&self, arguments: Value) -> ToolExecutionResult;
492
493    /// Execute the tool with context.
494    ///
495    /// This method provides access to runtime context like session ID and
496    /// optional stores (file store, etc.). Override this method for tools
497    /// that need access to session context or external resources.
498    ///
499    /// The default implementation simply calls `execute()`, ignoring the context.
500    ///
501    /// # Arguments
502    ///
503    /// * `arguments` - The arguments passed to the tool as a JSON value.
504    /// * `context` - Runtime context containing session ID and optional stores.
505    ///
506    /// # Returns
507    ///
508    /// A `ToolExecutionResult` indicating success, tool error, or internal error.
509    async fn execute_with_context(
510        &self,
511        arguments: Value,
512        _context: &ToolContext,
513    ) -> ToolExecutionResult {
514        // Default: delegate to execute(), ignoring context
515        self.execute(arguments).await
516    }
517
518    /// Returns true if this tool requires context for execution.
519    ///
520    /// Tools that need session context (like filesystem tools) should
521    /// override this to return true.
522    fn requires_context(&self) -> bool {
523        false
524    }
525
526    /// Returns the tool policy (auto or requires_approval).
527    ///
528    /// Default is `Auto` which means the tool executes immediately.
529    /// Override to return `RequiresApproval` for sensitive operations.
530    fn policy(&self) -> ToolPolicy {
531        ToolPolicy::Auto
532    }
533
534    /// Returns semantic hints describing the tool's behavioral properties.
535    ///
536    /// Override to provide hints like readonly, destructive, idempotent, etc.
537    /// Default is empty (all hints unspecified).
538    fn hints(&self) -> ToolHints {
539        ToolHints::default()
540    }
541
542    /// Returns backend-authored narration for a call to this tool, e.g.
543    /// "Read AGENTS.md".
544    ///
545    /// The owning capability's default [`crate::capabilities::Capability::narrate`]
546    /// dispatches here for the tool whose `name()` matches the call. Return
547    /// `None` to accept the generic `narration_noun`/display-name fallback.
548    /// Implementations should use the phrasing helpers in
549    /// [`crate::tool_narration`] (`narrate_read_file`, `narrate_shell_exec`, …)
550    /// so wording and localization stay consistent.
551    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    /// Returns native background execution support when this tool opts into
562    /// detached execution via `hints().supports_background`.
563    fn as_background_executable(&self) -> Option<&dyn BackgroundExecutableTool> {
564        None
565    }
566
567    /// Deferral policy for progressive tool-schema disclosure (tool search).
568    /// Hot-path or "consult-first" tools can return [`DeferrablePolicy::Never`]
569    /// to always keep their full schema directly callable. Defaults to
570    /// [`DeferrablePolicy::Automatic`].
571    fn deferrable_policy(&self) -> DeferrablePolicy {
572        DeferrablePolicy::default()
573    }
574
575    /// Convert this tool to a ToolDefinition for the agent config.
576    ///
577    /// This is used by ToolRegistry to generate tool definitions
578    /// for the LLM provider.
579    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// ============================================================================
595// ToolRegistry - Collection of Tools
596// ============================================================================
597
598/// A registry that holds multiple tools and implements ToolExecutor.
599///
600/// ToolRegistry provides a convenient way to manage multiple tools and
601/// integrate them with the agent loop. It implements `ToolExecutor` so
602/// it can be used directly with `AgentLoop`.
603///
604/// # Example
605///
606/// ```ignore
607/// use everruns_core::tools::{Tool, ToolRegistry};
608///
609/// // Create registry and add tools
610/// let mut registry = ToolRegistry::new();
611/// registry.register(Box::new(GetCurrentTime));
612/// registry.register(Box::new(GetWeather));
613///
614/// // Get tool definitions for agent config
615/// let definitions = registry.tool_definitions();
616///
617/// // Use with agent loop
618/// let agent_loop = AgentLoop::new(config, emitter, store, llm, registry);
619/// ```
620#[derive(Default, Clone)]
621pub struct ToolRegistry {
622    tools: HashMap<String, Arc<dyn Tool>>,
623}
624
625impl ToolRegistry {
626    /// Create a new empty tool registry
627    pub fn new() -> Self {
628        Self {
629            tools: HashMap::new(),
630        }
631    }
632
633    /// Create a tool registry with default built-in tools.
634    ///
635    /// This includes:
636    /// - `get_current_time`: Returns the current date and time
637    /// - `echo`: Echoes back the provided message
638    /// - `report_progress`: Emits deterministic external progress updates
639    /// - TestMath tools: add, subtract, multiply, divide
640    /// - TestWeather tools: get_weather, get_forecast
641    /// - TaskList tools: write_todos
642    /// - FileSystem tools: read_file, write_file, edit_file, list_directory, grep_files, delete_file, stat_file
643    /// - WebFetch tools: web_fetch — only when the `web-fetch` cargo feature is
644    ///   enabled (on by default; absent in provider builds that disable core
645    ///   default features)
646    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            // NOTE: `spawn_background` is intentionally NOT a default tool —
658            // it is contributed by the `background_execution` capability,
659            // which is auto-activated by
660            // `collect_capabilities_with_configs` whenever a collected tool
661            // declares `ToolHints::supports_background = Some(true)`. Keeping
662            // it out of defaults preserves the lockstep contract between
663            // model-visible tools and the worker execution registry: the
664            // executor only knows about `spawn_background` when the model
665            // can also see it.
666            .tool(ReportProgressTool)
667            // TestMath capability tools
668            .tool(AddTool)
669            .tool(SubtractTool)
670            .tool(MultiplyTool)
671            .tool(DivideTool)
672            // TestWeather capability tools
673            .tool(GetWeatherTool)
674            .tool(GetForecastTool)
675            // TaskList capability tools
676            .tool(WriteTodosTool)
677            // FileSystem capability tools
678            .tool(ReadFileTool)
679            .tool(WriteFileTool)
680            .tool(EditFileTool)
681            .tool(ListDirectoryTool)
682            .tool(GrepFilesTool)
683            .tool(DeleteFileTool)
684            .tool(StatFileTool);
685
686        // WebFetch capability tools (gated behind the `web-fetch` feature so
687        // provider crates that opt out of core defaults skip the fetchkit tree).
688        #[cfg(feature = "web-fetch")]
689        let builder = builder.tool(crate::capabilities::WebFetchTool::default());
690
691        builder.build()
692    }
693
694    /// Register a tool with the registry.
695    ///
696    /// If a tool with the same name already exists, it will be replaced.
697    pub fn register(&mut self, tool: impl Tool + 'static) {
698        self.tools.insert(tool.name().to_string(), Arc::new(tool));
699    }
700
701    /// Register a boxed tool
702    pub fn register_boxed(&mut self, tool: Box<dyn Tool>) {
703        self.tools.insert(tool.name().to_string(), Arc::from(tool));
704    }
705
706    /// Register an Arc-wrapped tool
707    pub fn register_arc(&mut self, tool: Arc<dyn Tool>) {
708        self.tools.insert(tool.name().to_string(), tool);
709    }
710
711    /// Get a tool by name
712    pub fn get(&self, name: &str) -> Option<&Arc<dyn Tool>> {
713        self.tools.get(name)
714    }
715
716    /// Check if a tool is registered
717    pub fn has(&self, name: &str) -> bool {
718        self.tools.contains_key(name)
719    }
720
721    /// Get the number of registered tools
722    pub fn len(&self) -> usize {
723        self.tools.len()
724    }
725
726    /// Check if the registry is empty
727    pub fn is_empty(&self) -> bool {
728        self.tools.is_empty()
729    }
730
731    /// Get all tool names
732    pub fn tool_names(&self) -> Vec<&str> {
733        self.tools.keys().map(|s| s.as_str()).collect()
734    }
735
736    /// Get tool definitions for use in RuntimeAgent.
737    ///
738    /// Returns a Vec of ToolDefinition that can be passed to
739    /// `RuntimeAgent::with_tools()`.
740    pub fn tool_definitions(&self) -> Vec<ToolDefinition> {
741        self.tools.values().map(|t| t.to_definition()).collect()
742    }
743
744    /// Remove a tool from the registry
745    pub fn unregister(&mut self, name: &str) -> Option<Arc<dyn Tool>> {
746        self.tools.remove(name)
747    }
748
749    /// Clear all tools from the registry
750    pub fn clear(&mut self) {
751        self.tools.clear();
752    }
753
754    /// Create a builder for fluent tool registration
755    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        // Use execute_with_context for all tools - context-aware tools will use it,
794        // regular tools will delegate to execute() via the default implementation
795        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
802// ============================================================================
803// ToolRegistryBuilder - Fluent API for Building Registry
804// ============================================================================
805
806/// Builder for creating a ToolRegistry with a fluent API.
807///
808/// # Example
809///
810/// ```ignore
811/// let registry = ToolRegistry::builder()
812///     .tool(GetCurrentTime)
813///     .tool(GetWeather)
814///     .build();
815/// ```
816pub struct ToolRegistryBuilder {
817    registry: ToolRegistry,
818}
819
820impl ToolRegistryBuilder {
821    /// Create a new builder
822    pub fn new() -> Self {
823        Self {
824            registry: ToolRegistry::new(),
825        }
826    }
827
828    /// Add a tool to the registry
829    pub fn tool(mut self, tool: impl Tool + 'static) -> Self {
830        self.registry.register(tool);
831        self
832    }
833
834    /// Add a boxed tool to the registry
835    pub fn tool_boxed(mut self, tool: Box<dyn Tool>) -> Self {
836        self.registry.register_boxed(tool);
837        self
838    }
839
840    /// Add an Arc-wrapped tool to the registry
841    pub fn tool_arc(mut self, tool: Arc<dyn Tool>) -> Self {
842        self.registry.register_arc(tool);
843        self
844    }
845
846    /// Build the registry
847    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
858// ============================================================================
859// Built-in Tools
860// ============================================================================
861
862/// A tool that echoes back its arguments (useful for testing)
863pub 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
912/// Spawn a background-capable tool and return immediately with a run handle.
913pub 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                    // Best-effort: create a monitor task linked to this schedule.
1161                    // The schedule is the source of truth; task creation failure
1162                    // does not fail the spawn.
1163                    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                                // Silent: the schedule's injected prompt message already
1185                                // wakes the session; a second wake on every fire would be noise.
1186                                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        // Create the session task tracking this run (specs/session-tasks.md).
1263        // task_registry is guaranteed Some above.
1264        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                    // Reaper uses this to decide whether to re-run on orphan.
1274                    // Only idempotent/readonly tools are safe to re-execute.
1275                    "reattachable": tool.hints().idempotent.unwrap_or(false)
1276                        || tool.hints().readonly.unwrap_or(false),
1277                    // Persisted so re-attach can restore the original signaling behavior.
1278                    "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        // Clone registry/ids for the cancel-watcher inside the spawned task.
1310        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        // Attempt captured at task creation for stale-attempt fencing.
1314        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            // Run the tool future inside a select! against a cancel-watch loop
1322            // when a task registry and task_id are wired up.  The watcher sends
1323            // a heartbeat every ~2 s and checks cancel_requested_at; when set,
1324            // it wins the select and the tool future is dropped.
1325            //
1326            // Rationale: cancel intent is recorded in shared storage so this
1327            // design works even when cancel_task executes on a different worker.
1328            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                            // Send heartbeat with stale-attempt fencing so a
1352                            // superseded executor's heartbeats are rejected once
1353                            // the reaper increments the attempt counter.
1354                            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                            // Check for cooperative cancel intent.
1366                            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                            // Cancel was requested; the tool future is dropped.
1378                            Err(ToolExecutionResult::ToolError(
1379                                BACKGROUND_CANCEL_SENTINEL.to_string(),
1380                            ))
1381                        }
1382                    }
1383                }
1384                // No registry or no task_id: run without cancel watch (unchanged behaviour).
1385                _ => 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            // Route canceled outcome through finalize_canceled so file/resource
1399            // cleanup is consistent with the success/fail paths.
1400            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    /// Session task mirroring this run; None when no task registry is wired.
1454    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    /// Mirror an update onto the session task (best-effort).
1487    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
1733/// Internal sentinel injected by the cancel-watch select branch. Namespaced
1734/// so a background tool that legitimately fails with "canceled" is never
1735/// misclassified as a cooperative cancel.
1736const BACKGROUND_CANCEL_SENTINEL: &str = "__everruns_background_cancel__";
1737
1738/// Returns true when an outcome from the cancel-watch select is the sentinel
1739/// cancel signal.  Factored out so the condition is testable without a
1740/// running async runtime.
1741fn 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
1747/// Re-attach a `background_tool` task after worker loss.
1748///
1749/// Called by `BackgroundToolTaskExecutor::start()` when the reaper decides the
1750/// task is safe to restart. Reads `spec["tool"]` and `spec["arguments"]` from
1751/// the task, looks up the tool in the built-in default registry, and spawns a
1752/// fresh background run with `task.attempt` as the heartbeat fence. New artifact
1753/// paths are generated so old partial artifacts do not conflict.
1754///
1755/// Returns an error (→ reaper falls back to orphaned-fail) when:
1756/// - `context.file_store` or `context.session_task_registry` is absent
1757/// - `spec["tool"]` is absent or empty
1758/// - the tool is not in the built-in default registry
1759/// - the tool does not implement `BackgroundExecutable`
1760/// - the tool's current hints are not `idempotent` or `readonly`
1761/// - per-worker or per-session background concurrency caps are exhausted
1762pub(crate) async fn reattach_background_run(
1763    task: &crate::session_task::SessionTask,
1764    context: &crate::traits::ToolContext,
1765) -> crate::error::Result<()> {
1766    // Fail fast before spawning a tokio task so the reaper can fall back to
1767    // orphaned-fail rather than leaving the task stuck in Running forever.
1768    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    // Re-verify tool hints from the live registry rather than trusting
1812    // spec["reattachable"], which could be forged via task creation APIs.
1813    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    // Enforce the same concurrency caps as spawn_background so many concurrent
1821    // re-attaches cannot exhaust worker or session limits.
1822    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    // Restore original signaling behavior; default true for tasks created before
1837    // this field was persisted.
1838    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        // Hold permits for the duration of the re-attached run.
1870        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
1961/// A tool that always fails (useful for testing error handling)
1962pub struct FailingTool {
1963    error_message: String,
1964    use_internal_error: bool,
1965}
1966
1967impl FailingTool {
1968    /// Create a failing tool with a tool-level error
1969    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    /// Create a failing tool with an internal error
1977    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// ============================================================================
2029// Tests
2030// ============================================================================
2031
2032#[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            // Test store is single-org; count all enabled schedules.
2718            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        // Success
2822        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        // Tool error (packaged as {"error": "..."} in result field, also sets error)
2828        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        // Internal error (packaged as {"error": "..."} with generic message)
2837        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        // GetCurrentTimeTool has display_name "Get Current Time"
2877        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        // Single-key object with _raw_output_scalar must not be mistaken for a
2935        // success_with_raw_output carrier when raw_output is absent.
2936        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        // Core tools
3003        assert!(
3004            registry.has("get_current_time"),
3005            "should have get_current_time"
3006        );
3007        assert!(registry.has("echo"), "should have echo");
3008        // spawn_background is contributed by the background_execution
3009        // capability (auto-activated) — it must NOT be in defaults.
3010        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        // TestMath capability tools
3021        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        // TestWeather capability tools
3027        assert!(registry.has("get_weather"), "should have get_weather");
3028        assert!(registry.has("get_forecast"), "should have get_forecast");
3029
3030        // TaskList capability tools
3031        assert!(registry.has("write_todos"), "should have write_todos");
3032
3033        // FileSystem capability tools
3034        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        // WebFetch capability tools
3043        assert!(registry.has("web_fetch"), "should have web_fetch");
3044
3045        // Total count: 19 - 1 (spawn_background, moved to capability) = 18
3046        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        // Test echo tool execution
3054        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        // Test add tool
3072        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        // AddTool returns floats, so compare as f64
3083        assert_eq!(result.result.unwrap()["result"].as_f64().unwrap(), 8.0);
3084    }
3085
3086    /// Regression: with_defaults() must NOT include capability-provided tools like
3087    /// 'bash'. These tools come from capabilities and must be registered separately.
3088    /// If bash were in defaults, the harness capability fallback would be masked.
3089    #[test]
3090    fn test_with_defaults_excludes_capability_only_tools() {
3091        let registry = ToolRegistry::with_defaults();
3092
3093        // bash comes from bashkit_shell capability, not defaults
3094        assert!(
3095            !registry.has("bash"),
3096            "bash must not be in defaults — it comes from bashkit_shell capability"
3097        );
3098        // kv_store/secret_store come from session_storage capability
3099        assert!(
3100            !registry.has("kv_store"),
3101            "kv_store must not be in defaults — it comes from session_storage capability"
3102        );
3103        // spawn_background comes from background_execution capability and is
3104        // auto-activated by `collect_capabilities_with_configs` when a
3105        // background-capable tool is present (see EVE-501).
3106        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; // still available in result json
3160
3161        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        // Wait for all tasks to be running (semaphore acquired).
3286        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        // The permit drop is enqueued in the spawned task after it marks the
3340        // resource Completed, so the cache entry may still exist briefly once
3341        // we observe Completed status.  Poll until pruned rather than asserting
3342        // immediately to avoid a race.
3343        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        // No task_registry wired — should fail.
3366        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        // No file_store wired — should fail.
3396        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    // =========================================================================
3587    // Cooperative cancellation tests
3588    // =========================================================================
3589
3590    #[test]
3591    fn test_is_canceled_outcome_detects_sentinel() {
3592        // The sentinel produced by the cancel-watcher branch.
3593        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    /// A background tool that sleeps indefinitely, allowing the test to exercise
3615    /// cancel via the cancel-watcher without actually waiting forever.
3616    #[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            // Sleep for a very long time; the cancel-watcher will win the select.
3631            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    // Minimal in-memory SessionTaskRegistry for cancel tests.
3675    // (Mirrors the double in capabilities/session_tasks.rs — kept local because
3676    //  that module is private.)
3677    #[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    /// End-to-end cancel test: spawn a long-sleeping background tool, then call
3786    /// `request_cancel` on the task registry, and assert the task ends Canceled.
3787    #[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        // Wait until the background task is Running (heartbeat loop started).
3821        tokio::time::timeout(std::time::Duration::from_secs(5), async {
3822            loop {
3823                // Wait for a heartbeat to confirm the watcher is live.
3824                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        // Request cancel.
3836        task_registry
3837            .request_cancel(session_id, &task_id)
3838            .await
3839            .expect("request_cancel should succeed");
3840
3841        // Wait for the task to reach Canceled state.
3842        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        // Verify result.json and output.log were written.
3856        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    // -------------------------------------------------------------------------
3881    // reattach_background_run early-guard tests
3882    // -------------------------------------------------------------------------
3883
3884    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        // Context with no file_store — only session_task_registry is wired.
3918        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        // Context has a file_store but no session_task_registry.
3942        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        // "test_background" is not in ToolRegistry::with_defaults().
3969        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        // Spec has no "tool" field.
3994        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}