aidaemon 0.11.1

use super::types::AbortOnDrop;
use crate::agent::*;
use crate::execution_policy::PolicyBundle;
use regex::RegexBuilder;
use serde_json::{json, Value};

pub(super) struct ToolExecutionIoResult {
    pub result_text: String,
    pub tool_duration_ms: u64,
    pub result_metadata: crate::traits::ToolCallMetadata,
}

pub(super) struct ToolExecutionIoCtx<'a> {
    pub effective_arguments: &'a str,
    pub idempotency_key: Option<&'a str>,
    pub injected_project_dir: Option<&'a str>,
    pub project_scope: Option<&'a str>,
    pub session_id: &'a str,
    pub task_id: &'a str,
    pub status_tx: &'a Option<mpsc::Sender<StatusUpdate>>,
    pub channel_ctx: &'a ChannelContext,
    pub user_role: UserRole,
    pub heartbeat: &'a Option<Arc<AtomicU64>>,
    pub emitter: &'a crate::events::EventEmitter,
    pub policy_bundle: &'a PolicyBundle,
}

pub(super) async fn execute_tool_call_io(
    agent: &Agent,
    tc: &ToolCall,
    ctx: &ToolExecutionIoCtx<'_>,
) -> ToolExecutionIoResult {
    let (start_label, start_summary) = crate::tools::sanitize::user_facing_tool_activity(
        &tc.name,
        &summarize_tool_args(&tc.name, ctx.effective_arguments),
    );
    send_status(
        ctx.status_tx,
        StatusUpdate::ToolStart {
            name: start_label,
            summary: start_summary,
        },
    );

    // Emit ToolCall event
    let _ = ctx
        .emitter
        .emit(
            EventType::ToolCall,
            ToolCallData::from_tool_call(
                tc.id.clone(),
                tc.name.clone(),
                serde_json::from_str(ctx.effective_arguments).unwrap_or(serde_json::json!({})),
                Some(ctx.task_id.to_string()),
            )
            .with_policy_metadata(
                ctx.idempotency_key
                    .map(str::to_string)
                    .or_else(|| Some(format!("{}:{}:{}", ctx.task_id, tc.name, tc.id))),
                Some(ctx.policy_bundle.policy.policy_rev),
                Some(ctx.policy_bundle.risk_score),
            ),
        )
        .await;

    let tool_exec_start = Instant::now();
    touch_heartbeat(ctx.heartbeat);

    // For long-running tools (cli_agent, terminal), spawn a background
    // task that keeps the heartbeat alive so the channel-level stale
    // watchdog doesn't auto-cancel the task while the tool is still
    // actively working.
    // Wrap in AbortOnDrop so the keeper is automatically cancelled if
    // handle_message is dropped by an outer select! (e.g. stale watchdog).
    // Without this, a detached keeper loop continues touching the heartbeat
    // forever, preventing the typing indicator's stale check from firing.
    let _heartbeat_keeper = if matches!(tc.name.as_str(), "cli_agent" | "terminal" | "spawn_agent")
    {
        ctx.heartbeat.as_ref().map(|hb| {
            let hb = Arc::clone(hb);
            AbortOnDrop(tokio::spawn(async move {
                loop {
                    tokio::time::sleep(Duration::from_secs(30)).await;
                    let now = SystemTime::now()
                        .duration_since(UNIX_EPOCH)
                        .unwrap_or_default()
                        .as_secs();
                    hb.store(now, Ordering::Relaxed);
                }
            }))
        })
    } else {
        None
    };

    let result = agent
        .execute_tool_with_watchdog_outcome(
            &tc.name,
            ctx.effective_arguments,
            &tool_exec::ToolExecCtx {
                session_id: ctx.session_id,
                task_id: Some(ctx.task_id),
                status_tx: ctx.status_tx.clone(),
                channel_visibility: ctx.channel_ctx.visibility,
                channel_id: ctx.channel_ctx.channel_id.as_deref(),
                project_scope: ctx.project_scope,
                trusted: ctx.channel_ctx.trusted,
                user_role: ctx.user_role,
            },
        )
        .await;

    // _heartbeat_keeper is dropped here (or when the scope ends),
    // which triggers AbortOnDrop to cancel the background task.
    drop(_heartbeat_keeper);
    touch_heartbeat(ctx.heartbeat);
    let mut result_metadata = crate::traits::ToolCallMetadata::default();
    let mut result_is_err = result.is_err();
    let mut result_text = match result {
        Ok(outcome) => {
            result_metadata = outcome.metadata;
            let text = outcome.output;
            // Sanitize and wrap untrusted tool outputs
            if !crate::tools::sanitize::is_trusted_tool(&tc.name) {
                let sanitized = crate::tools::sanitize::sanitize_external_content(&text);
                crate::tools::sanitize::wrap_untrusted_output(&tc.name, &sanitized)
            } else if tc.name == "terminal" {
                crate::tools::sanitize::strip_internal_control_markers(&text)
            } else {
                text
            }
        }
        Err(e) => {
            result_metadata.transport_error = Some(e.to_string());
            format!("Error: {}", e)
        }
    };

    if result_is_err && tc.name == "edit_file" {
        if let Some(recovered_text) =
            maybe_retry_edit_file_not_found_recovery(agent, &tc.arguments, &result_text, ctx).await
        {
            result_text = recovered_text;
            result_is_err = false;
            result_metadata.transport_error = None;
        }
    }

    if let Some(injected_dir) = ctx.injected_project_dir {
        result_text = format!(
            "{}\n\n{}",
            result_text,
            ToolResultNotice::PathAutoInjectedFromProjectContext {
                injected_dir: injected_dir.to_string(),
            }
            .render()
        );
    }

    // `cli_agent` errors can be extremely large (process output, stack traces).
    // Truncate aggressively to prevent context blow-up and runaway retries.
    if tc.name == "cli_agent" && result_is_err {
        let char_len = result_text.chars().count();
        if char_len > 2000 {
            let head: String = result_text.chars().take(500).collect();
            let tail: String = result_text
                .chars()
                .rev()
                .take(500)
                .collect::<Vec<char>>()
                .into_iter()
                .rev()
                .collect();
            result_text = format!(
                "{}\n\n[... cli_agent error output truncated ({} chars total) ...]\n\n{}",
                head, char_len, tail
            );
        }
    }

    // Compress large tool results to save context budget
    if agent.context_window_config.enabled {
        result_text = crate::memory::context_window::compress_tool_result(
            &tc.name,
            &result_text,
            agent.context_window_config.max_tool_result_chars,
        );
    }

    let tool_duration_ms = tool_exec_start.elapsed().as_millis().min(u64::MAX as u128) as u64;
    ToolExecutionIoResult {
        result_text,
        tool_duration_ms,
        result_metadata,
    }
}

async fn maybe_retry_edit_file_not_found_recovery(
    agent: &Agent,
    arguments: &str,
    initial_error: &str,
    ctx: &ToolExecutionIoCtx<'_>,
) -> Option<String> {
    if !initial_error.contains("Text not found in ") {
        return None;
    }

    let args: Value = serde_json::from_str(arguments).ok()?;
    let path = args.get("path")?.as_str()?.to_string();
    let old_text = args.get("old_text")?.as_str()?.to_string();
    if old_text.trim().is_empty() {
        return None;
    }

    let exec_ctx = tool_exec::ToolExecCtx {
        session_id: ctx.session_id,
        task_id: Some(ctx.task_id),
        status_tx: ctx.status_tx.clone(),
        channel_visibility: ctx.channel_ctx.visibility,
        channel_id: ctx.channel_ctx.channel_id.as_deref(),
        project_scope: ctx.project_scope,
        trusted: ctx.channel_ctx.trusted,
        user_role: ctx.user_role,
    };

    // Deterministic self-recovery path:
    // 1) Read current file state.
    // 2) Attempt one whitespace-tolerant mapping from old_text to exact on-disk text.
    // 3) Retry edit_file once with exact recovered old_text.
    let read_args = json!({ "path": path }).to_string();
    let read_probe_ok = agent
        .execute_tool_with_watchdog("read_file", &read_args, &exec_ctx)
        .await
        .is_ok();

    let resolved_path = crate::tools::fs_utils::validate_path(&path).ok()?;
    let file_content = tokio::fs::read_to_string(&resolved_path).await.ok()?;
    let recovered_old_text = recover_old_text_with_whitespace_tolerance(&file_content, &old_text)?;

    if recovered_old_text == old_text {
        return None;
    }

    let mut retry_args = args;
    retry_args["old_text"] = Value::String(recovered_old_text);
    let retry_args_str = serde_json::to_string(&retry_args).ok()?;
    match agent
        .execute_tool_with_watchdog("edit_file", &retry_args_str, &exec_ctx)
        .await
    {
        Ok(retry_output) => {
            let read_note = if read_probe_ok {
                "read_file probe succeeded"
            } else {
                "read_file probe failed, but direct file read succeeded"
            };
            Some(format!(
                "{}\n\n{}",
                retry_output,
                ToolResultNotice::InternalEditFileRecoverySucceeded {
                    read_note: read_note.to_string(),
                }
                .render()
            ))
        }
        Err(e) => {
            warn!(
                path = %path,
                error = %e,
                "Internal edit_file recovery retry failed"
            );
            None
        }
    }
}

fn build_whitespace_tolerant_pattern(old_text: &str) -> Option<String> {
    let mut pattern = String::new();
    let mut has_non_whitespace = false;
    let mut in_ws = false;

    for ch in old_text.chars() {
        if ch.is_whitespace() {
            if !in_ws {
                pattern.push_str(r"\s+");
                in_ws = true;
            }
        } else {
            has_non_whitespace = true;
            in_ws = false;
            pattern.push_str(&regex::escape(&ch.to_string()));
        }
    }

    if has_non_whitespace {
        Some(pattern)
    } else {
        None
    }
}

fn recover_old_text_with_whitespace_tolerance(content: &str, old_text: &str) -> Option<String> {
    let pattern = build_whitespace_tolerant_pattern(old_text)?;
    let regex = RegexBuilder::new(&pattern)
        .dot_matches_new_line(true)
        .build()
        .ok()?;

    let mut matches = regex.find_iter(content);
    let first = matches.next()?;
    if matches.next().is_some() {
        return None;
    }
    Some(content[first.start()..first.end()].to_string())
}

#[cfg(test)]
mod tests {
    use super::{build_whitespace_tolerant_pattern, recover_old_text_with_whitespace_tolerance};

    #[test]
    fn whitespace_tolerant_pattern_collapses_runs() {
        let pattern = build_whitespace_tolerant_pattern("foo   bar\tbaz\nqux").unwrap();
        assert_eq!(pattern, "foo\\s+bar\\s+baz\\s+qux");
    }

    #[test]
    fn recover_old_text_with_indentation_mismatch() {
        let content = "<section>\n    <h1>Dog World</h1>\n</section>\n";
        let old_text = "<section>\n  <h1>Dog World</h1>\n</section>\n";
        let recovered = recover_old_text_with_whitespace_tolerance(content, old_text).unwrap();
        assert_eq!(recovered, "<section>\n    <h1>Dog World</h1>\n</section>\n");
    }

    #[test]
    fn recover_old_text_returns_none_when_ambiguous() {
        let content = "alpha beta\nalpha    beta\n";
        let old_text = "alpha beta";
        assert!(recover_old_text_with_whitespace_tolerance(content, old_text).is_none());
    }
}