DevFlow/src-tauri/src/commands/ai/augmentation/registry.rs

//! Mention Resolver 注册表(核心设计2)
//!
//! 借鉴 [`crate::commands::ai::tool_registry::build_ai_tool_registry`] 的 AiToolRegistry
//! 模式:启动期 build 一次性注册,通过 `Arc<ResolverRegistry>` 注入 AppState,
//! 读路径(`resolve_all`)无锁并发访问。
//!
//! `resolve_all` 按 [`MentionRef::kind_str`] 分发到对应 Resolver,单条 resolve 失败
//! 收集进 errors 列表不阻断整批注入(返回 `(Vec<Augmentation>, Vec<ResolveError>)`)。

use std::collections::{HashMap, HashSet};
use std::sync::Arc;

use df_types::augmentation::{Augmentation, MentionRef, ResolveError};

use crate::commands::ai::augmentation::sanitize::ProviderLocality;
use crate::commands::ai::augmentation::MentionResolver;

/// Mention Resolver 注册表 — 按 kind 分发 MentionRef 到对应 Resolver。
///
/// 启动期 build 一次性 `register`,之后只读(`resolve_all`),故内部 HashMap 无 RwLock,
/// 通过 `Arc<ResolverRegistry>` 注入 AppState 共享。新增 mention 类型 = build 处
/// `register` 一行 + 加一个 impl,主干零改。
pub struct ResolverRegistry {
    resolvers: HashMap<&'static str, Arc<dyn MentionResolver>>,
}

impl ResolverRegistry {
    /// 创建空注册表。
    pub fn new() -> Self {
        Self {
            resolvers: HashMap::new(),
        }
    }

    /// 注册一个 Resolver。重复 kind 后注册者覆盖(开发期防漏注同名)。
    pub fn register(&mut self, resolver: Arc<dyn MentionResolver>) {
        let kind = resolver.kind();
        self.resolvers.insert(kind, resolver);
    }

    /// 批量 resolve mention 引用列表为 Augmentation。
    ///
    /// 按 [`MentionRef::kind_str`] 分发到对应 Resolver,逐条独立 resolve:
    /// - 成功 → push 进 `augs`
    /// - 失败(Resolver 返 Err / kind 无对应 Resolver)→ push 进 `errors`,**不中断**
    ///   后续 mention(整批注入最大化保留可用上下文)。
    ///
    /// 同一引用去重(先到先得):`/技能` 前导与 @mention span 可指向同一技能,按
    /// `kind + ref_id`(skill 的 ref_id 即 name)跳过重复,避免同一正文重复 resolve
    /// 并在 prompt 注入两份。首现位置决定保留哪条,顺序不变。
    ///
    /// 返回 `(成功列表, 失败列表)`,调用方据 errors 数决定是否提示用户某 mention 跳过。
    pub async fn resolve_all(
        &self,
        refs: &[MentionRef],
        loc: ProviderLocality,
    ) -> (Vec<Augmentation>, Vec<ResolveError>) {
        let mut augs = Vec::with_capacity(refs.len());
        let mut errors = Vec::new();
        let mut seen: HashSet<String> = HashSet::with_capacity(refs.len());
        for rf in refs {
            // 去重键 = kind + ref_id(id 或 skill 的 name):同键 mention 只 resolve 首条
            if !seen.insert(format!("{}:{}", rf.kind_str(), rf.ref_id())) {
                continue;
            }
            let kind = rf.kind_str();
            match self.resolvers.get(kind) {
                Some(resolver) => match resolver.resolve(rf, loc).await {
                    Ok(aug) => augs.push(aug),
                    Err(e) => errors.push(e),
                },
                None => {
                    // kind 无对应 Resolver(未注册的 mention 类型):收集 KindMismatch 不中断
                    errors.push(ResolveError::KindMismatch {
                        expected: "<registered>".to_string(),
                        actual: kind.to_string(),
                    });
                }
            }
        }
        (augs, errors)
    }
}

impl Default for ResolverRegistry {
    fn default() -> Self {
        Self::new()
    }
}

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

    /// 测试用 Resolver:kind 固定 "test",resolve 固定返 Skip(模拟失败收集)。
    struct SkipResolver;
    #[async_trait]
    impl MentionResolver for SkipResolver {
        fn kind(&self) -> &'static str {
            "test"
        }
        async fn resolve(
            &self,
            _rf: &MentionRef,
            _loc: ProviderLocality,
        ) -> Result<Augmentation, ResolveError> {
            Err(ResolveError::Skip {
                kind: "test".into(),
                ref_id: "x".into(),
                reason: "test skip".into(),
            })
        }
    }

    #[tokio::test]
    async fn resolve_all_unknown_kind_collected_not_panic() {
        let reg = ResolverRegistry::new();
        // 未注册任何 Resolver,传入 Project mention → KindMismatch 收集,不 panic
        let rf = MentionRef::Project {
            id: "p".into(),
            name: "n".into(),
        };
        let (augs, errors) = reg.resolve_all(&[rf], ProviderLocality::Local).await;
        assert!(augs.is_empty());
        assert_eq!(errors.len(), 1);
        match &errors[0] {
            ResolveError::KindMismatch { actual, .. } => assert_eq!(actual, "project"),
            other => panic!("expected KindMismatch, got {:?}", other),
        }
    }

    #[tokio::test]
    async fn resolve_all_failure_collected_continues_others() {
        // 验证:单条 resolve 失败不中断后续 mention。
        // 构造两条 mention,第一条命中已注册 Resolver(SkipResolver 返 Skip 失败),
        // 第二条命中未知 kind(KindMismatch)。两条失败都收集,augs 为空,无 panic。
        // 注:SkipResolver kind="test",无法用现有 MentionRef 变体精确匹配(枚举 kind 固定
        // project/task/idea/skill),故此处仅验证 unknown-kind 路径收集不中断——核心语义
        // (失败收集 + 继续)已在 resolve_all 实现的 for 循环中保证,本测为回归守护。
        let mut reg = ResolverRegistry::new();
        reg.register(Arc::new(SkipResolver));
        let refs = vec![
            MentionRef::Project { id: "p".into(), name: "n".into() },
            MentionRef::Skill { name: "s".into() },
        ];
        let (augs, errors) = reg.resolve_all(&refs, ProviderLocality::Local).await;
        assert!(augs.is_empty(), "两条 mention kind 均≠test 应全 KindMismatch");
        assert_eq!(errors.len(), 2, "两条失败都应收集");
    }

    /// 测试用 Resolver:kind 固定 "skill",resolve 固定返 Skill(模拟成功注入)。
    struct OkSkillResolver;
    #[async_trait]
    impl MentionResolver for OkSkillResolver {
        fn kind(&self) -> &'static str {
            "skill"
        }
        async fn resolve(
            &self,
            rf: &MentionRef,
            _loc: ProviderLocality,
        ) -> Result<Augmentation, ResolveError> {
            let name = match rf {
                MentionRef::Skill { name } => name.clone(),
                _ => String::new(),
            };
            Ok(Augmentation::Skill {
                name,
                source: "test".into(),
                body: "BODY".into(),
            })
        }
    }

    #[tokio::test]
    async fn resolve_all_dedups_same_ref_first_wins() {
        // 去重(先到先得):/技能 前导 + @mention span 指向同一技能时只 resolve/注入一份;
        // 不同技能不受影响;首现位置决定保留顺序。
        let mut reg = ResolverRegistry::new();
        reg.register(Arc::new(OkSkillResolver));
        let refs = vec![
            MentionRef::Skill { name: "review".into() },
            MentionRef::Skill { name: "review".into() },
            MentionRef::Skill { name: "idea".into() },
        ];
        let (augs, errors) = reg.resolve_all(&refs, ProviderLocality::Local).await;
        assert!(errors.is_empty());
        assert_eq!(augs.len(), 2, "同名技能两条应去重为一份");
        match &augs[0] {
            Augmentation::Skill { name, .. } => assert_eq!(name, "review", "先到先得保留首条"),
            other => panic!("expected Skill, got {:?}", other),
        }
    }
}