DevFlow/src-tauri/src/commands/ai/tool_failure_memory.rs

//! B5d 工具失败记忆(conv 级):治「404 连环十几次无退避换路」——AI 工具是原生能力,
//! 应有学习能力:同 (conv, tool, args) 连败 ≥2 次 → 拦截重试,回填结构化失败提示,
//! 引导 LLM 换路(改参数/换工具/报告放弃);成功即清零(知识保鲜,防误杀)。
//!
//! 设计要点(与既有机制的关系):
//! - 与 fail_window/熔断(stall_detector)互补:熔断治「会话级停滞」(连续 N 轮全失败),
//!   本模块治「单点死磕」(同一调用原样重发);两层独立计数互不干扰。
//! - 与只读缓存(find_cached_readonly_result)不同:缓存复用成功结果,本模块拦截失败重复。
//! - 记忆生命周期 = conv 生命周期(进程内 HashMap,conv 结束随会话对象丢弃;不落库,
//!   下次会话重新学习——失败参数往往任务态相关,跨会话记忆弊大于利)。
//! - key 归一:tool_name + args 关键字段(read_file 用 path,grep 用 pattern+path,
//!   run_command 用命令行本体;其余工具全 args JSON)——LLM 微调无关字段(如 limit)
//!   不应逃脱记忆,但完全相同调用必须命中。

use std::collections::HashMap;
use std::sync::{Mutex, OnceLock};

/// 连败拦截阈值:同 key 连败 ≥2 次(即第 3 次尝试)在执行前拦截。
pub const ABORT_AFTER_CONSECUTIVE_FAILURES: u32 = 2;

/// 记忆条目:连续失败次数 + 最近一次错误摘要(拦截时回填给 LLM 换路提示)。
#[derive(Debug, Clone)]
struct FailureEntry {
    consecutive: u32,
    last_error: String,
}

/// 进程级失败记忆(全 conv 共享单锁;锁内零 await,微秒级)。
static FAILURE_MEMORY: OnceLock<Mutex<HashMap<String, FailureEntry>>> = OnceLock::new();

fn memory() -> &'static Mutex<HashMap<String, FailureEntry>> {
    FAILURE_MEMORY.get_or_init(|| Mutex::new(HashMap::new()))
}

/// args → 记忆 key 的归一化提取(工具感知)。
///
/// 原则:抽「语义主键」而非全 args——read_file 微调 limit 重读同一文件,应视为同调用;
/// run_command 命令行本体即主键(含路径/参数)。
fn memory_key(tool: &str, args: &serde_json::Value) -> String {
    let get = |k: &str| args.get(k).and_then(|v| v.as_str()).unwrap_or("").to_string();
    match tool {
        "read_file" | "read_symbol" | "file_info" | "delete_file" => {
            format!("{tool}|path={}", get("path"))
        }
        "grep" => format!("{tool}|pat={}|path={}", get("pattern"), get("path")),
        "search_files" => format!("{tool}|pat={}|path={}", get("pattern"), get("path")),
        "list_directory" => format!("{tool}|path={}", get("path")),
        "run_command" => format!("{tool}|cmd={}", get("command")),
        "mysql_query" | "mysql_execute" => format!("{tool}|sql={}", get("sql")),
        _ => {
            // 通用:全 args 稳定序列化(serde_json Value 保序不可保证 → BTreeMap 排序归一)
            let sorted = serde_json::to_string(&sort_json(args)).unwrap_or_default();
            format!("{tool}|{sorted}")
        }
    }
}

/// serde_json Value 递归排序(Object 按 key 排序),保证 key 稳定。
fn sort_json(v: &serde_json::Value) -> serde_json::Value {
    use serde_json::Value;
    match v {
        Value::Object(m) => {
            let mut sorted: Vec<(String, Value)> =
                m.iter().map(|(k, v)| (k.clone(), sort_json(v))).collect();
            sorted.sort_by(|a, b| a.0.cmp(&b.0));
            Value::Object(sorted.into_iter().collect())
        }
        Value::Array(a) => Value::Array(a.iter().map(sort_json).collect()),
        other => other.clone(),
    }
}

/// 执行前检查:同 key 连败达标 → 拦截,返回结构化失败提示(LLM 可读换路引导)。
/// 未达标/无记忆 → None 放行。
pub fn precheck(tool: &str, args: &serde_json::Value) -> Option<String> {
    let key = memory_key(tool, args);
    let mem = memory().lock().unwrap_or_else(|p| p.into_inner());
    let Some(entry) = mem.get(&key) else { return None };
    if entry.consecutive >= ABORT_AFTER_CONSECUTIVE_FAILURES {
        Some(format!(
            "{{\"status\":\"blocked_by_failure_memory\",\"tool\":\"{tool}\",\"consecutive_failures\":{},\"last_error\":\"{}\",\"hint\":\"同一调用已连败 {} 次,已自动拦截。请勿原样重试:①检查参数(路径/SQL/命令拼写) ②换工具或换路径 ③若确认无法绕过,向用户报告放弃原因。\"}}",
            entry.consecutive,
            entry.last_error.replace('\"', "'"),
            entry.consecutive
        ))
    } else {
        None
    }
}

/// 执行后回写:成功清零删除,失败递增并记最近错误。
pub fn record(tool: &str, args: &serde_json::Value, ok: bool, error: Option<String>) {
    let key = memory_key(tool, args);
    let mut mem = memory().lock().unwrap_or_else(|p| p.into_inner());
    if ok {
        mem.remove(&key);
    } else {
        let e = mem.entry(key).or_insert(FailureEntry { consecutive: 0, last_error: String::new() });
        e.consecutive += 1;
        e.last_error = error.unwrap_or_default().chars().take(200).collect();
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    fn clear() {
        memory().lock().unwrap().clear();
    }

    #[test]
    fn precheck_blocks_after_two_failures() {
        clear();
        let args = json!({"path": "/tmp/x.rs"});
        assert!(precheck("read_file", &args).is_none());
        record("read_file", &args, false, Some("文件不存在".into()));
        assert!(precheck("read_file", &args).is_none()); // 1 次未达阈值
        record("read_file", &args, false, Some("文件不存在".into()));
        let blocked = precheck("read_file", &args).unwrap();
        assert!(blocked.contains("blocked_by_failure_memory"));
        assert!(blocked.contains("consecutive_failures\":2"));
        // 无关字段微调不逃脱记忆:limit 改了仍拦截(同 path 主键)
        let args2 = json!({"path": "/tmp/x.rs", "limit": 300});
        assert!(precheck("read_file", &args2).is_some());
        // 成功清零
        record("read_file", &args, true, None);
        assert!(precheck("read_file", &args).is_none());
    }

    #[test]
    fn memory_key_normalizes_irrelevant_fields() {
        let a = json!({"path": "/a.rs", "limit": 10});
        let b = json!({"path": "/a.rs", "limit": 999});
        assert_eq!(memory_key("read_file", &a), memory_key("read_file", &b));
        // 通用工具:全 args 排序归一,键序无关
        let c = json!({"a": 1, "b": 2});
        let d = serde_json::from_str::<serde_json::Value>(r#"{"b":2,"a":1}"#).unwrap();
        assert_eq!(memory_key("generic_tool", &c), memory_key("generic_tool", &d));
    }

    #[test]
    fn run_command_key_is_command_itself() {
        let a = json!({"command": "cargo test", "working_dir": "E:/x"});
        assert!(memory_key("run_command", &a).contains("cmd=cargo test"));
    }
}