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