zesdex/src/app/subagent/engine.rs

//! Subagent execution loop: drive an LLM conversation, gate tool calls
//! against the context's allowlist, run tools, and stream progress events
//! to the parent via an mpsc channel.
//!
//! Security: subagent tool gating mirrors the main agent's `Harness` checks
//! (path traversal, reason validation, stub/denial/assumption scanning,
//! bash exfiltration and destructive-pattern detection) so that subagents
//! are not a weaker link than the main agent.

use std::fmt::Write;
use tokio::sync::mpsc;
use crate::dto::chat::message::ChatMessage;
use crate::dto::provider::request::ToolDef;
use crate::tool::{all_tools, tool_defs, tool_is_risky};
use super::context::SubagentContext;
use super::event::SubagentEvent;

/// Maps a subagent's allowed tool names to concrete Tool trait objects and
/// OpenAI-style tool definitions.
///
/// Flow: load `all_tools()` → if `allowed_tools` is empty, use all; else
/// filter by membership → derive `ToolDef`s for the LLM.
///
/// Why: an empty allowlist means "no restriction" (matches
/// `build_subagent_context`'s default for non-reviewer roles).
///
/// Return: `(tool impls, schema defs)` for the subagent to use.
fn build_subagent_tools(allowed_tools: &[String]) -> (Vec<Box<dyn crate::tool::Tool>>, Vec<ToolDef>) {
    let all = all_tools();
    let filtered: Vec<Box<dyn crate::tool::Tool>> = if allowed_tools.is_empty() {
        all.into_iter()
            .filter(|t| t.name() != "hive_mind" && t.name() != "workflow_run")
            .collect()
    } else {
        all.into_iter()
            .filter(|t| {
                allowed_tools.contains(&t.name().to_string())
                    && t.name() != "hive_mind"
                    && t.name() != "workflow_run"
            })
            .collect()
    };
    let defs = tool_defs(&filtered);
    (filtered, defs)
}

/// Resolve the API key, model, and base URL from persisted app config.
///
/// Flow: try the settings key for the active provider → fall back to the
/// provider's `api_key_env` env-var → fall back to the provider's
/// `default_api_key` → fall back to an empty string.
///
/// Why: matches the main agent's credential resolution exactly, so
/// subagents automatically inherit the same provider settings.
///
/// Return: `(api_key, model, optional_base_url)`.
fn resolve_provider_config() -> (String, String, Option<String>) {
    let settings = crate::model::settings::Settings::load();
    let app_config = crate::model::app_config::AppConfig::load();

    let mut api_key = settings.api_keys.get(&settings.provider).cloned().unwrap_or_else(|| {
        tracing::warn!("[subagent] no API key for provider '{}' in settings, trying env/default", settings.provider);
        String::new()
    });
    let model = settings.model.clone();
    let base_url = app_config.providers.get(&settings.provider)
        .map(|p| p.api_base.clone());

    if api_key.is_empty() {
        if let Some(provider_cfg) = app_config.providers.get(&settings.provider) {
            api_key = provider_cfg.api_key_env.as_ref()
                .and_then(|env| std::env::var(env).ok())
                .or_else(|| provider_cfg.default_api_key.clone())
                .unwrap_or_else(|| {
                    tracing::warn!("[subagent] all API key resolution paths exhausted for '{}'", settings.provider);
                    String::new()
                });
        }
    }

    (api_key, model, base_url)
}

// ─── Subagent-level tool gating (mirrors Harness checks) ───

const STUB_PATTERNS: &[&str] = &[
    "todo!()", "todo!(",
    "unimplemented!()", "unimplemented!(",
    "FIXME", "fixme:", "XXX:", "PLACEHOLDER",
    "REPLACE_ME", "stub_value", "stub_function",
    "fake_response", "fake_data",
    "not implemented", "not yet implemented",
    "to be implemented", "to be done",
];

const DENIAL_PATTERNS: &[&str] = &[
    "// skip", "// skipping", "// skipping for now",
    "// for now just", "// punt", "// hack:",
    "// workaround:", "// cba", "// later",
    "// do later", "// ignore for now", "// disable",
    "// bypass", "// quick fix", "// temp fix",
    "// temporary fix", "// temp:", "// temporary:",
    "// noop",
];

const ASSUMPTION_PATTERNS: &[&str] = &[
    "// assume", "// probably", "// guess",
    "// should work", "// hopefully", "// i think",
    "// should be fine", "// likely",
];

const EXFIL_PATTERNS: &[&str] = &[
    "curl ", "wget ", "nc -e ", "ncat ", "/dev/tcp/",
    "base64 -d |", "base64 --decode |",
    "openssl s_client", "ssh -R ",
    "scp /", "rsync /",
];

const SENSITIVE_PATH_PATTERNS: &[&str] = &[
    ".ssh/id_rsa", ".ssh/id_ed25519",
    ".aws/credentials", ".aws/config",
    ".kube/config", ".docker/config.json",
    "/etc/shadow", "/etc/passwd", "/proc/self/environ",
];

const MIN_REASON_LEN: usize = 8;

/// Gate a tool call in the subagent context. Returns `Some(block_reason)` if
/// the call should be blocked, `None` to allow.
///
/// Flow: always blocks dangerous patterns — path traversal, stub/denial/
/// assumption language, bash exfiltration, destructive commands, sensitive
/// path reads — regardless of the allowed-tools list. Tools that are not
/// risky only get the basic allowlist check.
fn gate_subagent_tool_call(
    tool_name: &str,
    args: &serde_json::Value,
) -> Option<String> {
    // File-mutating tools: write / edit / delete
    if matches!(tool_name, "write" | "edit" | "delete") {
        if let Some(path) = args.get("path").and_then(|v| v.as_str()) {
            if path.contains("..") {
                return Some("path traversal detected in 'path' argument".to_string());
            }
        }
    }

    // write / edit require a non-trivial `reason`
    if matches!(tool_name, "write" | "edit" | "delete") {
        let reason = args.get("reason").and_then(|v| v.as_str()).unwrap_or("");
        if reason.trim().len() < MIN_REASON_LEN {
            return Some(format!(
                "{tool_name} requires a non-trivial 'reason' (>= {MIN_REASON_LEN} chars) explaining why",
            ));
        }
    }

    // write / edit content must not contain stubs, denial, or assumption language
    if matches!(tool_name, "write" | "edit") {
        let content = match tool_name {
            "write" => args.get("content").and_then(|v| v.as_str()).unwrap_or(""),
            "edit" => {
                let old = args.get("old").and_then(|v| v.as_str()).unwrap_or("");
                let new = args.get("new").and_then(|v| v.as_str()).unwrap_or("");
                // For edits, scanning old+new together catches stubs in both
                return if contains_any(old, STUB_PATTERNS) || contains_any(new, STUB_PATTERNS) {
                    Some("content contains stub/placeholder pattern; production code must be fully implemented".to_string())
                } else if contains_any(new, DENIAL_PATTERNS) {
                    Some("content contains denial/punt pattern; implement properly instead of skipping".to_string())
                } else if contains_any(new, ASSUMPTION_PATTERNS) {
                    Some("content contains assumption pattern; verify against data instead of guessing".to_string())
                } else {
                    return None;
                };
            }
            _ => "",
        };
        if contains_any(content, STUB_PATTERNS) {
            return Some("content contains stub/placeholder pattern; production code must be fully implemented".to_string());
        }
        if contains_any(content, DENIAL_PATTERNS) {
            return Some("content contains denial/punt pattern; implement properly instead of skipping".to_string());
        }
        if contains_any(content, ASSUMPTION_PATTERNS) {
            return Some("content contains assumption pattern; verify against data instead of guessing".to_string());
        }
    }

    // Bash: exfiltration, sensitive paths, destructive commands
    if tool_name == "bash" {
        let cmd = args.get("command").and_then(|v| v.as_str()).unwrap_or("");
        if cmd.contains("..") {
            return Some("path traversal detected in bash command".to_string());
        }
        // Only check exfiltration for non-standard commands
        let is_standard = cmd.trim_start().starts_with("cargo")
            || cmd.trim_start().starts_with("rustc")
            || cmd.trim_start().starts_with("git ")
            || cmd.trim_start().starts_with("ls")
            || cmd.trim_start().starts_with("pwd")
            || cmd.trim_start().starts_with("echo")
            || cmd.trim_start().starts_with("cat")
            || cmd.trim_start().starts_with("find")
            || cmd.trim_start().starts_with("grep")
            || cmd.trim_start().starts_with("test");
        if !is_standard {
            for pat in EXFIL_PATTERNS {
                if cmd.contains(pat) {
                    return Some(format!("potential data-exfiltration command blocked (matched '{pat}')"));
                }
            }
        }
        for pat in SENSITIVE_PATH_PATTERNS {
            if cmd.contains(pat) {
                return Some(format!("refused to read/write sensitive path '{pat}'"));
            }
        }
        let dangerous = ["rm -rf /", "rm -rf --no-preserve-root", "rm -rf ~",
            "rm -fr /", "mkfs.", "dd if=", ":(){", "> /dev/sda",
            "chmod -R 000 /", "shutdown ", "poweroff ", "reboot ", "halt "];
        for pat in &dangerous {
            if cmd.contains(pat) {
                return Some(format!("destructive command pattern blocked: {pat}"));
            }
        }
        if contains_any(cmd, STUB_PATTERNS) {
            return Some("bash command contains stub pattern".to_string());
        }
    }

    // git_operator: require reason
    if tool_name == "git_operator" {
        let reason = args.get("reason").and_then(|v| v.as_str()).unwrap_or("");
        if reason.trim().len() < MIN_REASON_LEN {
            return Some("git_operator requires a non-trivial 'reason' (>= 8 chars)".to_string());
        }
    }

    None
}

/// Check if `text` matches any pattern (case-insensitive substring).
fn contains_any(text: &str, patterns: &[&str]) -> bool {
    let lower = text.to_lowercase();
    patterns.iter().any(|p| lower.contains(&p.to_lowercase()))
}

/// Build an ASCII tree of the workspace directory structure for the
/// system prompt, so the LLM can see the file layout.
///
/// Flow: for each root, walk using `ignore::WalkBuilder` (respecting
/// `.gitignore` and hidden files) → prefix `[DIR]` for directories →
/// truncate after 1000 entries.
fn generate_workspace_tree(roots: &[std::path::PathBuf]) -> String {
    let mut out = String::new();
    out.push_str("Current Workspace Directory Structure:\n");
    for root in roots {
        writeln!(out, "Root: {}", root.display()).unwrap();
        let walker = ignore::WalkBuilder::new(root)
            .hidden(true)
            .git_ignore(true)
            .build();
        let mut count = 0;
        for entry in walker.flatten() {
            let path = entry.path();
            if let Ok(rel) = path.strip_prefix(root) {
                if rel.as_os_str().is_empty() { continue; }
                let is_dir = entry.file_type().is_some_and(|ft| ft.is_dir());
                let prefix = if is_dir { "[DIR] " } else { "      " };
                writeln!(out, "  {}{}", prefix, rel.display()).unwrap();
                count += 1;
                if count > 1000 {
                    out.push_str("  ... (truncated)\n");
                    break;
                }
            }
        }
    }
    out
}

/// Synchronous subagent entry point: run up to `ctx.max_steps` iterations
/// of the LLM tool loop.
///
/// Flow: inject system prompt (with workspace tree if available) → for each
/// step: resolve provider config, build an LLM client, call
/// `chat_with_tools_streaming` (with abort check per SSE event), process
/// tool calls (gated against both the allowlist and Harness-style content
/// safety checks) or collect text output → send `SubagentEvent`s on `tx` →
/// break on first text-only (non-empty) response.
///
/// Why: runs synchronously on a dedicated thread so the main async event
/// loop is not blocked. Tool gating prevents restricted, risky, or
/// malicious/poor-quality tool calls from executing.
///
/// Return: the concatenated text output, or an `anyhow::Error` if the LLM
/// call fails at any step.
#[allow(clippy::too_many_lines)]
pub fn run_subagent(ctx: &SubagentContext, tx: &mpsc::Sender<SubagentEvent>) -> anyhow::Result<String> {
    let mut output = String::new();
    let mut messages: Vec<ChatMessage> = Vec::new();

    // Build system prompt with workspace tree context if we have workspaces,
    // giving subagents the same project-awareness as the main agent.
    let system_with_context = if ctx.workspaces.is_empty() {
        ctx.system_prompt.clone()
    } else {
        let tree_info = generate_workspace_tree(&ctx.workspaces);
        format!("{}\n\n{}", ctx.system_prompt, tree_info)
    };
    messages.push(ChatMessage::system(system_with_context));

    let tool_ctx = crate::tool::ToolCtx::builder()
        .session_dir(ctx.session_dir.clone())
        .workspaces(ctx.workspaces.clone())
        .origin(crate::app::state::types::Origin::SubAgent)
        .workflow_findings(ctx.workflow_findings.clone())
        .build();

    // Build tool list once before the loop
    let (tools, tdefs) = build_subagent_tools(&ctx.allowed_tools);
    let tdefs_opt: Option<Vec<ToolDef>> = if tdefs.is_empty() { None } else { Some(tdefs) };

    // Cache provider config once before the loop instead of re-resolving
    // from disk on every step (Settings::load + AppConfig::load each parse
    // JSON files, and the config cannot change between steps).
    let (api_key, model, base_url) = resolve_provider_config();
    let client = crate::service::provider::LlmClient::new(api_key, model, base_url);

    for step in 0..ctx.max_steps {

        // Check abort flag before each LLM call so a stuck subagent can
        // be cancelled from the parent (mirrors main agent behaviour).
        if ctx.abort_flag.as_ref().is_some_and(|f| f.load(std::sync::atomic::Ordering::SeqCst)) {
            let _ = tx.blocking_send(SubagentEvent::StepFailed {
                step,
                error: "subagent aborted by parent".to_string(),
            });
            anyhow::bail!("subagent aborted by parent at step {step}");
        }

        // Use streaming API so the abort flag is checked per SSE event,
        // making the subagent responsive to cancellation even during an
        // LLM call (non-streaming would block for 10-30s unchecked).
        let stream_result = client.chat_with_tools_streaming(
            &messages,
            tdefs_opt.clone(),
            Some(0.7),
            Some(4096),
            |_event| -> bool {
                // Check abort on every SSE event for responsive cancellation.
                if ctx.abort_flag.as_ref().is_some_and(|f| f.load(std::sync::atomic::Ordering::SeqCst)) {
                    return false; // signals provider to abort
                }
                // We don't stream tokens to the UI for subagents — just
                // need the assembled message at the end.
                true
            },
        );

        let (response, _usage) = match stream_result {
            Ok(result) => result,
            Err(e) => {
                let is_abort = ctx.abort_flag.as_ref().is_some_and(|f| f.load(std::sync::atomic::Ordering::SeqCst))
                    || e.to_string().contains("aborted");
                let _ = tx.blocking_send(SubagentEvent::StepFailed {
                    step,
                    error: if is_abort {
                        "subagent aborted by user".to_string()
                    } else {
                        e.to_string()
                    },
                });
                if is_abort {
                    anyhow::bail!("subagent aborted by parent at step {step}");
                }
                // No non-streaming fallback — API must support streaming.
                // Non-streaming calls block for up to 1 min without checking
                // abort_flag, making cancellation unresponsive.
                anyhow::bail!("subagent call failed at step {step}: {e}");
            }
        };

        let has_tool_calls = response.tool_calls.is_some()
            && response.tool_calls.as_ref().is_some_and(|tc| !tc.is_empty());

        let content = response.content.clone().unwrap_or_default();

        if has_tool_calls {
            let tool_calls = response.tool_calls.clone().unwrap_or_default();
            // Push the assistant message with tool_calls into the conversation
            messages.push(response);

            let mut results_vec = Vec::new();
            std::thread::scope(|s| {
                let mut handles = Vec::new();
                let tools_ref = &tools;
                let tool_ctx_ref = &tool_ctx;
                for tool_call in &tool_calls {
                    let handle = s.spawn(move || {
                        // Check abort flag before each tool execution
                        if ctx.abort_flag.as_ref().is_some_and(|f| f.load(std::sync::atomic::Ordering::SeqCst)) {
                            return (tool_call, Err(anyhow::anyhow!("subagent aborted by parent during tool execution")));
                        }

                        let tool_name = &tool_call.function.name;
                        let args = crate::dto::chat::tool::sanitize_tool_arguments(&tool_call.function.arguments);
                        let explicitly_allowed = ctx.allowed_tools.contains(tool_name);
                        let generally_allowed = ctx.allowed_tools.is_empty() || explicitly_allowed;

                        // Level 1: allowlist check — is this tool even permitted?
                        if !generally_allowed {
                            return (tool_call, Ok(format!("tool '{tool_name}' not allowed for this subagent")));
                        }

                        // Level 2: risky tool check — risky tools require explicit permission
                        if tool_is_risky(tool_name) && !explicitly_allowed {
                            return (tool_call, Ok(format!("risky tool '{tool_name}' requires explicit permission; not allowed for this subagent")));
                        }

                        // Level 3: Harness-style content safety gating
                        if let Some(block_reason) = gate_subagent_tool_call(tool_name, &args) {
                            return (tool_call, Ok(format!("Blocked by subagent gate: {block_reason}")));
                        }

                        let result = match tools_ref.iter().find(|t| t.name() == tool_name.as_str()) {
                            Some(tool) => {
                                let is_edit = tool_name == "write" || tool_name == "edit";
                                if is_edit && !tool_call.id.is_empty() {
                                    if let Ok(conn) = crate::model::msglog::open_or_create(&ctx.session_dir) {
                                        let path = args.get("path").and_then(|v| v.as_str()).unwrap_or("");
                                        if let Ok(abs_path) = crate::tool::resolve_path(&tool_ctx_ref.workspaces, path) {
                                            if let Ok(bytes) = std::fs::read(&abs_path) {
                                                let session_id = ctx.session_dir
                                                    .file_name()
                                                    .and_then(|n| n.to_str())
                                                    .unwrap_or("unknown");
                                                let _ = crate::model::msglog::store_blob(
                                                    &conn, session_id, &tool_call.id, &bytes, None,
                                                );
                                            }
                                        }
                                    }
                                }

                                let run_res = tool.run(tool_ctx_ref, &args);

                                if is_edit && run_res.is_ok() {
                                    let reason = args
                                        .get("reason")
                                        .and_then(|v| v.as_str())
                                        .unwrap_or("unnamed");
                                    let path = args
                                        .get("path")
                                        .and_then(|v| v.as_str())
                                        .unwrap_or("unknown");
                                    let content_sha256 = {
                                        let content = args.get("content").or_else(|| args.get("new"));
                                        use sha2::Digest;
                                        let hash = sha2::Sha256::digest(
                                            content.and_then(|v| v.as_str()).unwrap_or("").as_bytes(),
                                        );
                                        hex::encode(hash)
                                    };
                                    let bytes_delta = if tool_name == "write" {
                                        args.get("content")
                                            .and_then(|v| v.as_str())
                                            .map_or(0, |s| s.len() as i64)
                                    } else {
                                        let old = args.get("old").and_then(|v| v.as_str()).unwrap_or("");
                                        let new = args.get("new").and_then(|v| v.as_str()).unwrap_or("");
                                        (new.len() as i64 - old.len() as i64).abs()
                                    };
                                    let session_id = ctx.session_dir
                                        .file_name()
                                        .and_then(|n| n.to_str())
                                        .unwrap_or("unknown")
                                        .to_string();
                                    let entry = crate::model::editlog::EditLogEntry {
                                        ts: chrono::Utc::now().timestamp_millis(),
                                        tool: tool_name.clone(),
                                        path: path.to_string(),
                                        reason: reason.to_string(),
                                        content_sha256,
                                        bytes_delta,
                                        origin: tool_ctx_ref.origin.tag(),
                                        session_id,
                                    };
                                    let mut el = crate::model::editlog::EditLog::new(&ctx.session_dir);
                                    el.append(entry).ok();
                                }
                                run_res
                            }
                            None => Err(anyhow::anyhow!("tool '{tool_name}' not found")),
                        };
                        (tool_call, result)
                    });
                    handles.push(handle);
                }
                for h in handles {
                    if let Ok(res) = h.join() {
                        results_vec.push(res);
                    }
                }
            });

            for (tool_call, result) in results_vec {
                let tool_name = &tool_call.function.name;
                let args = crate::dto::chat::tool::sanitize_tool_arguments(&tool_call.function.arguments);

                let _ = tx.blocking_send(SubagentEvent::ToolCall {
                    tool: tool_name.clone(),
                    args: args.clone(),
                });

                match result {
                    Ok(output_text) => {
                        messages.push(ChatMessage::tool_result(tool_call.id.clone(), output_text.clone()));
                        let _ = tx.blocking_send(SubagentEvent::ToolResult {
                            tool: tool_name.clone(),
                            output: output_text,
                        });
                    }
                    Err(e) => {
                        let err_str = e.to_string();
                        if err_str.contains("subagent aborted by parent") {
                            let _ = tx.blocking_send(SubagentEvent::StepFailed {
                                step,
                                error: err_str.clone(),
                            });
                            anyhow::bail!("{err_str}");
                        }
                        let msg = format!("tool '{tool_name}' failed: {e}");
                        messages.push(ChatMessage::tool_result(tool_call.id.clone(), msg.clone()));
                        let _ = tx.blocking_send(SubagentEvent::ToolResult {
                            tool: tool_name.clone(),
                            output: msg,
                        });
                    }
                }
            }
        } else {
            // Text-only response — accumulate and finish
            if !content.is_empty() {
                output.push_str(&content);
                output.push('\n');
            }
            let _ = tx.blocking_send(SubagentEvent::StepCompleted {
                step,
                output: content.clone(),
            });
            // Break only when we got real content; empty means something went wrong
            if !content.is_empty() {
                break;
            }
        }
    }

    let _ = tx.blocking_send(SubagentEvent::Completed { output: output.clone() });
    Ok(output)
}