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