diff --git a/src/memory/retrieval/fast.rs b/src/memory/retrieval/fast.rs index 2322b9a..318552e 100644 --- a/src/memory/retrieval/fast.rs +++ b/src/memory/retrieval/fast.rs @@ -208,7 +208,7 @@ fn resolve_local( let mut hits = Vec::new(); for (id, _) in ordered { if let Some(mut hit) = by_id.remove(&id) { - if source_scope.is_some_and(|scope| !scope.contains(&hit.tree_scope)) { + if source_scope.is_some_and(|scope| !source_scope_allows(scope, &hit.tree_scope)) { continue; } hit.score = coverage.get(&id).copied().unwrap_or_default(); @@ -239,7 +239,9 @@ async fn dense( ) .await?; if let Some(scope) = source_scope { - response.hits.retain(|hit| scope.contains(&hit.tree_scope)); + response + .hits + .retain(|hit| source_scope_allows(scope, &hit.tree_scope)); } let total = response.hits.len(); response.hits.truncate(limit); @@ -282,6 +284,19 @@ fn dedup_ids(ids: impl Iterator) -> Vec { ids.filter(|id| seen.insert(id.clone())).collect() } +fn source_scope_allows(scope: &HashSet, tree_scope: &str) -> bool { + if scope.contains(tree_scope) { + return true; + } + extract_mem_src_id(tree_scope).is_some_and(|id| scope.contains(id)) +} + +fn extract_mem_src_id(value: &str) -> Option<&str> { + let rest = value.strip_prefix("mem_src:")?; + let (id, _) = rest.split_once(':')?; + (!id.is_empty()).then_some(id) +} + #[cfg(test)] #[path = "fast_tests.rs"] mod tests; diff --git a/src/memory/retrieval/source.rs b/src/memory/retrieval/source.rs index fffeb78..077f00f 100644 --- a/src/memory/retrieval/source.rs +++ b/src/memory/retrieval/source.rs @@ -236,6 +236,9 @@ fn scope_matches_kind(scope: &str, kind_prefix: &str) -> bool { if lower.starts_with(&format!("{kind_prefix}:")) { return true; } + if kind_prefix == SourceKind::Document.as_str() && lower.starts_with("mem_src:") { + return true; + } PLATFORM_KINDS .iter() .any(|(platform, kind)| *kind == kind_prefix && lower.starts_with(&format!("{platform}:"))) diff --git a/src/memory/retrieval/source_tests.rs b/src/memory/retrieval/source_tests.rs index 343aded..b8d0e3d 100644 --- a/src/memory/retrieval/source_tests.rs +++ b/src/memory/retrieval/source_tests.rs @@ -195,6 +195,10 @@ fn scope_prefix_matching_known_platforms() { assert!(scope_matches_kind("gmail:alice", "email")); assert!(scope_matches_kind("notion:page123", "document")); assert!(scope_matches_kind("linear:conn-1:issue-abc", "document")); + assert!(scope_matches_kind( + "mem_src:src-folder-9:Slides_Notes/example.md", + "document" + )); assert!(!scope_matches_kind("slack:#eng", "email")); assert!(scope_matches_kind("chat:custom", "chat")); }