oxo-flow-cli 0.25.0

//! AI-powered error recovery for run failures.

use anyhow::Result;
use colored::Colorize;
use oxo_flow_ai::{knowledge::builtin, provider::AiProvider, session::archive_before_modify};
use std::path::{Path, PathBuf};

/// Diagnostic result from AI error analysis.
pub struct DiagnoseResult {
    pub root_cause: String,
    pub fix_action: String,
    pub modified_toml: Option<String>,
    pub safe_to_auto_apply: bool,
}

/// Parse the model's safety verdict from the mandated output line
/// `- Safe to auto-apply: <yes/no>`.
///
/// The previous substring test (`contains("safe to auto-apply")`) matched
/// every conforming reply — including `- Safe to auto-apply: no` — so the
/// auto-apply gate in `run.rs` never blocked anything. Anything other than an
/// explicit `yes`/`true` verdict (missing, unparseable, or an explicit
/// `NOT safe to auto-apply` line) means: do not rewrite the user's workflow.
fn parse_safe_to_auto_apply(response_text: &str) -> bool {
    // An explicit refusal wins wherever it appears in the reply.
    if response_text
        .to_ascii_lowercase()
        .contains("not safe to auto-apply")
    {
        return false;
    }
    for line in response_text.lines() {
        // `to_ascii_lowercase` preserves byte length, so the index stays valid
        // on the original line.
        let lower = line.to_ascii_lowercase();
        let Some(idx) = lower.find("safe to auto-apply") else {
            continue;
        };
        let verdict = line[idx + "safe to auto-apply".len()..]
            .trim_start_matches(|c: char| c == ':' || c == '*' || c == '`' || c.is_whitespace())
            .trim_matches(|c: char| c == '*' || c == '_' || c == '`' || c.is_whitespace())
            .to_ascii_lowercase();
        if verdict.starts_with("yes") || verdict.starts_with("true") {
            return true;
        }
    }
    false
}

/// Analyze a pipeline failure using AI.
pub async fn diagnose_failure(
    workflow_path: &Path,
    failed_rule: &str,
    exit_code: i32,
    stderr: &str,
    provider: &AiProvider,
) -> Result<DiagnoseResult> {
    eprintln!();
    eprintln!("{}", "AI Error Recovery".bold().green());
    eprintln!("  Failed rule: {}", failed_rule.yellow());
    eprintln!("  Exit code: {}", exit_code.to_string().red());
    eprintln!(
        "  Model: {}\n",
        provider.model().unwrap_or_else(|| "default".into())
    );

    let system = build_diagnose_prompt();
    let toml_content =
        std::fs::read_to_string(workflow_path).unwrap_or_else(|_| "(unable to read)".into());

    let user = format!(
        "## Pipeline Failure\n\n\
         Rule: **{failed_rule}**\n\
         Exit code: {exit_code}\n\n\
         ## Error Output\n```\n{stderr}\n```\n\n\
         ## Current Workflow\n```toml\n{toml_content}\n```\n\n\
         ## Task\n\
         1. Match the error against known patterns\n\
         2. Identify the root cause\n\
         3. Propose a specific fix (change exact lines in the TOML)\n\
         4. Output the corrected TOML if changes are needed (inside ```toml fences)\n\
         5. State whether the fix is safe to auto-apply",
    );

    eprintln!("{}", "  Diagnosing...".bold().cyan());

    use oxo_flow_ai::types::Message;
    let messages = vec![Message::system(&system), Message::user(&user)];
    let response = provider.chat_with_tools(&messages, &[]).await?;
    let response_text = response.content.unwrap_or_default();

    // Record token usage
    tracing::info!(
        tokens_in = response.usage.prompt_tokens,
        tokens_out = response.usage.completion_tokens,
        "AI recovery analysis completed"
    );

    // Parse the AI response
    let root_cause = extract_section(&response_text, "Root Cause");
    let fix_action = extract_section(&response_text, "Fix");
    let safe = parse_safe_to_auto_apply(&response_text);

    // Extract modified TOML
    let modified_toml = extract_toml_block(&response_text);

    eprintln!(
        "\n{}\n{}",
        "Root Cause:".bold().red(),
        root_cause
            .as_deref()
            .unwrap_or("Unknown — see full analysis below")
    );
    eprintln!(
        "{}\n{}",
        "Suggested Fix:".bold().yellow(),
        fix_action
            .as_deref()
            .unwrap_or("Manual review needed — see full analysis below")
    );
    eprintln!(
        "Safe to auto-apply: {}",
        if safe { "yes".green() } else { "no".red() }
    );
    // When section parsing failed the model answered off-format; surface the
    // raw analysis instead of a dead-end pointer (the old "see full analysis"
    // wording pointed at output that was never printed).
    if root_cause.is_none() || fix_action.is_none() {
        eprintln!(
            "\n{}\n{}",
            "Full Analysis:".bold().underline(),
            response_text
        );
    }

    Ok(DiagnoseResult {
        root_cause: root_cause.unwrap_or_else(|| "Unknown — see full analysis".into()),
        fix_action: fix_action.unwrap_or_else(|| "Manual review needed".into()),
        modified_toml,
        safe_to_auto_apply: safe,
    })
}

/// Apply a fix to the workflow file, archiving the original.
pub fn apply_fix(workflow_path: &Path, modified_toml: &str, session_id: &str) -> Result<PathBuf> {
    // Archive original via the shared #544/P3-10 machinery: the snapshot
    // carries the full pre-fix workflow content, so it must land 0600 and
    // the per-workflow archive dir must stay pruned — `archive_dir_for`
    // used to hand-roll the layout with neither.
    let original = std::fs::read_to_string(workflow_path)?;
    let backup_path = archive_before_modify(workflow_path, &original, session_id)?;
    eprintln!(
        "{} Original archived to {}",
        "  ✓".green(),
        backup_path.display()
    );

    // Write fix
    std::fs::write(workflow_path, modified_toml)?;
    eprintln!(
        "{} Fix applied to {}",
        "  ✓".green(),
        workflow_path.display()
    );

    Ok(backup_path)
}

fn build_diagnose_prompt() -> String {
    let error_patterns = builtin::format_error_patterns();

    format!(
        r#"## Role & Identity
You are a senior bioinformatics pipeline incident responder for oxo-flow. When a pipeline rule fails,
you diagnose the root cause with surgical precision and propose the minimal fix needed.
Your analysis must be specific, actionable, and safe.

## Error Pattern Reference (Common Failures)
{error_patterns}

## Diagnostic Protocol (Execute in Order)
### Step 1 — Triage
- Classify the failure: resource (OOM/timeout), data (missing/corrupt input), software (version/API change), configuration (wrong params/paths)
- Map exit code to category: 137=OOM, 139=segfault, 1=user error, 127=command not found

### Step 2 — Root Cause Analysis
- Read the error output CAREFULLY — the answer is usually in the last 5 lines
- Cross-reference against error patterns above
- Check: is the tool version compatible with the provided parameters?
- Check: are input files actually produced by the dependency rule?
- Check: does the working directory have sufficient disk space?

### Step 3 — Fix Proposal
- Identify the EXACT line(s) in the TOML that need changing
- Propose the MINIMAL change — do not restructure the workflow unless essential
- If the fix is a resource adjustment (threads/memory), specify EXACT new values
- If the fix requires a tool parameter change, verify against the tool's documentation
- NEVER propose removing QC steps, validation, or safety checks

### Step 4 — Safety Assessment
- **Safe to auto-apply**: parameter tuning (threads, memory), fixing a typo, adding a missing flag
- **NOT safe to auto-apply**: adding/removing rules, changing DAG edges, modifying shell logic, changing file paths

## Output Format
Respond in this exact structure:

```
## Root Cause
<1-2 sentences explaining WHY the failure occurred, citing specific evidence from the error output>

## Proposed Fix
<Specific action. If changing TOML parameters, list old→new values. Reference exact lines.>

## Corrected TOML
```toml
<complete corrected [workflow] section or [[rules]] block — only the changed portions>
```

## Safety Assessment
- Safe to auto-apply: <yes/no>
- Risk level: <low/medium/high>
- Rollback: <how to undo this change>

## Prevention
<1 sentence on how to prevent this failure in future pipelines>
```

⚠ ONLY propose changes that fix the reported error. Do NOT make unrelated improvements.
"#
    )
}

/// Extract TOML code block from AI response.
///
/// The diagnostic prompt allows "only the changed portions" — a `[[rules]]`
/// block is a valid fix on its own, so accept it alongside `[workflow]`.
fn extract_toml_block(response: &str) -> Option<String> {
    if let Some(start) = response.find("```toml") {
        let start = start + 7;
        if let Some(end) = response[start..].find("```") {
            let content = response[start..start + end].trim().to_string();
            if content.contains("[workflow]") || content.contains("[[rules]") {
                return Some(content);
            }
        }
    }
    None
}

/// Extract a named section from the AI response.
///
/// Tries the inline patterns first (`Marker:`, `**Marker**:`, ...), then the
/// markdown headings the diagnostic prompt itself mandates (`## Root Cause`,
/// `## Proposed Fix`). The heading form is what conforming replies actually
/// use — without it every well-formed answer fell through to the fallback
/// text and the real analysis was dropped.
fn extract_section(text: &str, marker: &str) -> Option<String> {
    let patterns = [
        format!("[{marker}]:"),
        format!("[{marker}]"),
        format!("**{marker}**:"),
        format!("{marker}:"),
    ];
    for pat in &patterns {
        if let Some(pos) = text.find(pat.as_str()) {
            let remainder = &text[pos + pat.len()..];
            let end = remainder
                .find("\n[")
                .or_else(|| remainder.find("\n**"))
                .or_else(|| remainder.find("\n```"))
                .unwrap_or(remainder.len());
            let found = remainder[..end].trim();
            if !found.is_empty() {
                return Some(found.to_string());
            }
        }
    }
    extract_heading_section(text, marker)
}

/// Capture the body under a markdown heading whose title contains `marker`
/// (case-insensitive), ending at the next heading of any level or a code fence.
fn extract_heading_section(text: &str, marker: &str) -> Option<String> {
    let marker_lower = marker.to_ascii_lowercase();
    for (idx, line) in text.lines().enumerate() {
        let trimmed = line.trim();
        if !trimmed.starts_with('#') {
            continue;
        }
        let title = trimmed.trim_start_matches('#').trim().to_ascii_lowercase();
        if !title.contains(&marker_lower) {
            continue;
        }
        let body: String = text
            .lines()
            .skip(idx + 1)
            .take_while(|l| {
                let t = l.trim();
                !(t.starts_with('#') || t.starts_with("```"))
            })
            .collect::<Vec<_>>()
            .join("\n");
        let body = body.trim();
        if !body.is_empty() {
            return Some(body.to_string());
        }
    }
    None
}

#[cfg(test)]
mod tests {
    use super::apply_fix;
    use super::{extract_section, extract_toml_block, parse_safe_to_auto_apply};
    use std::path::{Path, PathBuf};

    #[test]
    fn safe_gate_requires_an_explicit_yes() {
        assert!(parse_safe_to_auto_apply("- Safe to auto-apply: yes"));
        assert!(parse_safe_to_auto_apply("- **Safe to auto-apply**: YES"));
        assert!(parse_safe_to_auto_apply("Safe to auto-apply: true"));
        // The regression: the old substring test matched this exact line.
        assert!(!parse_safe_to_auto_apply("- Safe to auto-apply: no"));
        assert!(!parse_safe_to_auto_apply("- Safe to auto-apply: <yes/no>"));
        assert!(!parse_safe_to_auto_apply(
            "- **NOT safe to auto-apply**: changing DAG edges"
        ));
        // Missing/unparseable verdicts must never authorize a rewrite.
        assert!(!parse_safe_to_auto_apply(
            "## Root Cause\nbwa ran out of memory"
        ));
        assert!(!parse_safe_to_auto_apply(""));
    }

    /// A fully-conforming reply uses the heading format the prompt mandates.
    #[test]
    fn section_extraction_matches_prompt_mandated_headings() {
        let reply = "## Root Cause\nThe rule requested 16 threads but the \
machine exposes 4; the scheduler killed it (exit 137).\n\n## Proposed Fix\n\
Change `threads = 16` to `threads = 4` in rule `align`.\n\n## Corrected TOML\n\
```toml\n[[rules]]\nname = \"align\"\nthreads = 4\n```\n\n## Safety Assessment\n\
- Safe to auto-apply: yes\n- Risk level: low\n";

        let cause = extract_section(reply, "Root Cause").expect("root cause parsed");
        assert!(cause.contains("exit 137"), "got: {cause}");
        let fix = extract_section(reply, "Fix").expect("fix parsed");
        assert!(fix.contains("threads = 4"), "got: {fix}");
    }

    #[test]
    fn section_extraction_stops_at_next_heading_or_fence() {
        let reply = "## Root Cause\nbad reference genome\n\n## Proposed Fix\npoint to hg38";
        assert_eq!(
            extract_section(reply, "Root Cause").as_deref(),
            Some("bad reference genome")
        );
        assert_eq!(
            extract_section(reply, "Fix").as_deref(),
            Some("point to hg38")
        );
    }

    #[test]
    fn section_extraction_still_supports_inline_markers() {
        assert_eq!(
            extract_section("**Root Cause**: typo in path", "Root Cause").as_deref(),
            Some("typo in path")
        );
        assert_eq!(
            extract_section("Root Cause: typo in path", "Root Cause").as_deref(),
            Some("typo in path")
        );
    }

    /// The prompt permits "only the changed portions" — a [[rules]]-only
    /// block is a legitimate fix and must not be discarded.
    #[test]
    fn toml_block_accepts_rules_only_fragments() {
        let fragment = "```toml\n[[rules]]\nname = \"align\"\nthreads = 4\n```";
        assert!(extract_toml_block(fragment).is_some());
        assert!(extract_toml_block("```toml\n[workflow]\nname = \"x\"\n```").is_some());
        assert!(extract_toml_block("```toml\nnot toml at all\n```").is_none());
        assert!(extract_toml_block("no fence here").is_none());
    }

    /// Serializes the tests that `set_current_dir` into a temp cwd —
    /// `apply_fix` → `archive_before_modify` resolves the archive root
    /// against the process cwd, and a parallel test mid-dance would
    /// otherwise observe the wrong root. A poisoned lock (one test
    /// panicking while the cwd is elsewhere) must not cascade into every
    /// sibling failing with PoisonError.
    static CWD_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());

    /// `apply_fix` resolves the archive root from the process cwd, so these
    /// tests run inside a temp cwd (pattern from web `workspace.rs` tests).
    fn with_temp_cwd(tag: &str, f: impl FnOnce(&Path)) {
        let _cwd = CWD_LOCK.lock().unwrap_or_else(|e| e.into_inner());
        let tmp = tempfile::Builder::new()
            .prefix(&format!("oxo-recover-{tag}-"))
            .tempdir()
            .unwrap();
        let prev = std::env::current_dir().unwrap();
        std::env::set_current_dir(tmp.path()).unwrap();

        // A project-local `.oxo-flow` marker makes `resolve_oxo_flow_dir`
        // anchor here instead of falling back to the real `$HOME` — the
        // layout assertions below must never touch the developer's home.
        std::fs::create_dir_all(".oxo-flow").unwrap();
        std::fs::write("workflow.oxoflow", "[workflow]\nname = \"t\"\n").unwrap();

        struct RestoreCwd(PathBuf);
        impl Drop for RestoreCwd {
            fn drop(&mut self) {
                // Runs on panic too — a cwd left inside a deleted temp dir
                // would make every other test fail resolving relative paths.
                let _ = std::env::set_current_dir(&self.0);
            }
        }
        let _restore = RestoreCwd(prev);

        f(tmp.path());
    }

    #[test]
    fn recovered_archive_is_owner_only() {
        // Issue #964: recover-side backups carry the full pre-fix workflow
        // content — they must land 0600 exactly like the orchestrator-side
        // archives (`archived_snapshot_is_owner_only` in oxo-flow-ai).
        with_temp_cwd("perm", |_| {
            let backup = apply_fix(
                Path::new("workflow.oxoflow"),
                "[workflow]\nname = \"fixed\"\n",
                "sess-964",
            )
            .unwrap();

            #[cfg(unix)]
            {
                use std::os::unix::fs::PermissionsExt;
                assert_eq!(
                    std::fs::metadata(&backup).unwrap().permissions().mode() & 0o777,
                    0o600,
                    "recover archive must be owner-only"
                );
            }
            let _ = backup;
        });
    }

    #[test]
    fn recovered_archive_prunes_oldest_snapshots() {
        // Issue #964: repeated recover runs previously grew
        // `.oxo-flow/ai_archive/<stem>/` without bound; the shared machinery
        // caps it (200 kept here — verified via the prune contract on a
        // smaller keep by `archive_prune_keeps_the_newest_snapshots_per_workflow`,
        // so this test only pins that the cap applies on the recover path).
        with_temp_cwd("prune", |_| {
            // Pre-fill the per-workflow archive dir beyond the cap with
            // older-dated snapshots, then let apply_fix add one more.
            let dir = std::env::current_dir()
                .unwrap()
                .join(".oxo-flow/ai_archive/workflow");
            std::fs::create_dir_all(&dir).unwrap();
            for i in 0..200 {
                let p = dir.join(format!("20250101-0000{i:02}-old{i}-before.oxoflow"));
                std::fs::write(&p, "old").unwrap();
                // Distinct mtimes: a same-second batch would leave the sort
                // order to the filesystem.
                let t = filetime::FileTime::from_unix_time(1_700_000_000 + i as i64, 0);
                filetime::set_file_mtime(&p, t).unwrap();
            }

            apply_fix(
                Path::new("workflow.oxoflow"),
                "[workflow]\nname = \"fixed\"\n",
                "sess-prune",
            )
            .unwrap();

            let remaining: Vec<_> = std::fs::read_dir(&dir)
                .unwrap()
                .filter_map(|e| e.ok())
                .map(|e| e.file_name().to_string_lossy().into_owned())
                .collect();
            assert_eq!(
                remaining.len(),
                200,
                "recover archives must be pruned to the cap: got {}",
                remaining.len()
            );
            assert!(
                !remaining.iter().any(|n| n.contains("old0-")),
                "the oldest pre-filled snapshot must be pruned: {remaining:?}"
            );
            assert!(
                remaining.iter().any(|n| n.contains("sess-prune")),
                "the fresh snapshot must survive the prune: {remaining:?}"
            );
        });
    }

    #[test]
    fn recovered_archive_layout_matches_session_machinery() {
        // The recover path must not retype its own archive layout: the
        // backup lives under `.oxo-flow/ai_archive/<stem>/` with the
        // `{timestamp}-{session}-before.oxoflow` name the shared helper
        // owns (and the fix is actually written).
        with_temp_cwd("layout", |_| {
            let backup = apply_fix(
                Path::new("workflow.oxoflow"),
                "[workflow]\nname = \"fixed\"\n",
                "sess-layout",
            )
            .unwrap();

            let expected_parent = std::env::current_dir()
                .unwrap()
                .join(".oxo-flow/ai_archive/workflow");
            assert_eq!(
                backup.parent().unwrap(),
                expected_parent,
                "archive dir must come from the shared resolver"
            );
            assert!(
                backup
                    .file_name()
                    .unwrap()
                    .to_string_lossy()
                    .ends_with("-sess-layout-before.oxoflow"),
                "backup filename must follow the shared convention: {}",
                backup.display()
            );
            assert_eq!(
                std::fs::read_to_string("workflow.oxoflow").unwrap(),
                "[workflow]\nname = \"fixed\"\n",
                "the fix itself must be written"
            );
        });
    }
}