//! Mention Resolver 注册表(核心设计2) //! //! 借鉴 [`crate::commands::ai::tool_registry::build_ai_tool_registry`] 的 AiToolRegistry //! 模式:启动期 build 一次性注册,通过 `Arc` 注入 AppState, //! 读路径(`resolve_all`)无锁并发访问。 //! //! `resolve_all` 按 [`MentionRef::kind_str`] 分发到对应 Resolver,单条 resolve 失败 //! 收集进 errors 列表不阻断整批注入(返回 `(Vec, Vec)`)。 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` 注入 AppState 共享。新增 mention 类型 = build 处 /// `register` 一行 + 加一个 impl,主干零改。 pub struct ResolverRegistry { resolvers: HashMap<&'static str, Arc>, } impl ResolverRegistry { /// 创建空注册表。 pub fn new() -> Self { Self { resolvers: HashMap::new(), } } /// 注册一个 Resolver。重复 kind 后注册者覆盖(开发期防漏注同名)。 pub fn register(&mut self, resolver: Arc) { 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, Vec) { let mut augs = Vec::with_capacity(refs.len()); let mut errors = Vec::new(); let mut seen: HashSet = 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: "".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 { 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 { 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), } } }