use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use super::external_origin::ChunkOrigin;
use super::{BM25Index, ChunkOrigins, CodeChunk, SearchResult};
use crate::core::context_admission::stores::StoreAdmission;
use crate::core::providers::provenance::{ProviderOrigin, with_reuse_deadline_until};
pub(super) const SEARCH_TOOL: &str = "ctx_search";
const MAX_LINEAGE_CHECKS: usize = 16;
pub(crate) const MAX_SOURCES: usize = 128;
const MAX_PROVIDER_CHECKS: usize = 16;
const PROVIDER_BUDGET: std::time::Duration = std::time::Duration::from_secs(3);
pub struct SourceGate {
root: PathBuf,
tool: &'static str,
sources: HashMap<String, Option<String>>,
authorities: HashMap<String, bool>,
admission: StoreAdmission,
origins: Arc<ChunkOrigins>,
provider_checks: usize,
lineage_checks: usize,
provider_deadline: Option<std::time::Instant>,
scope_root: PathBuf,
protected: Option<bool>,
withheld: usize,
}
impl SourceGate {
pub(crate) fn new(root: &Path, tool: &'static str) -> Self {
Self::with_admission(
root,
tool,
&crate::core::context_admission::carrier::gateway_snapshot(),
)
}
pub(crate) fn with_admission(
root: &Path,
tool: &'static str,
admission: &StoreAdmission,
) -> Self {
Self {
root: root.to_path_buf(),
tool,
sources: HashMap::new(),
authorities: HashMap::new(),
admission: admission.clone(),
origins: Arc::default(),
provider_checks: 0,
lineage_checks: 0,
provider_deadline: None,
scope_root: root.canonicalize().unwrap_or_else(|_| root.to_path_buf()),
protected: None,
withheld: 0,
}
}
fn in_root_scope<T>(&self, operation: impl FnOnce() -> T) -> T {
use crate::core::policy::runtime::REQUEST_PROJECT;
let bound = REQUEST_PROJECT
.try_with(|slot| slot.borrow().clone())
.ok()
.flatten();
if bound.is_some_and(|bound| {
bound == self.scope_root || bound.canonicalize().is_ok_and(|b| b == self.scope_root)
}) {
return operation();
}
REQUEST_PROJECT.sync_scope(
std::cell::RefCell::new(Some(self.scope_root.clone())),
operation,
)
}
pub(crate) fn protected(&mut self) -> bool {
if let Some(protected) = self.protected {
return protected;
}
let protected = self.in_root_scope(|| {
crate::core::context_admission::stored::source_authority_required_under(
self.admission.admission_policy(),
)
});
self.protected = Some(protected);
protected
}
#[must_use]
pub(crate) fn with_origins(mut self, origins: &Arc<ChunkOrigins>) -> Self {
self.origins = Arc::clone(origins);
self
}
fn over_bound(&self) -> bool {
self.sources.len() + self.authorities.len() >= MAX_SOURCES
}
fn authority_holds(&mut self, source: &str, check: impl FnOnce(&Self) -> bool) -> bool {
if let Some(&holds) = self.authorities.get(source) {
return holds;
}
if self.over_bound() {
return false;
}
let holds = self.in_root_scope(|| check(self));
self.authorities.insert(source.to_owned(), holds);
holds
}
fn admits_verbatim(&self, text: &str, label: Option<&str>) -> bool {
use crate::core::context_admission as admission;
self.in_root_scope(|| {
let admitted = match label {
Some(label) => admission::admit_source_with_snapshot(
text,
label,
false,
admission::OutputCompleteness::complete(),
&self.admission,
)
.ok(),
None => admission::carrier::admit_carrier_text_under(
text,
Some(self.admission.admission_policy()),
&admission::carrier::STORED,
)
.ok(),
};
admitted.is_some_and(|admitted| admitted.storable && admitted.text == text)
})
}
fn originless(&mut self, text: &str) -> bool {
!self.protected() && self.admits_verbatim(text, None)
}
fn served_without_file_lineage(&mut self) -> bool {
if self.protected() {
crate::core::archive::authority::mark_unbound();
}
true
}
pub(super) fn decided(&self, source: &str, text: &str) -> Option<bool> {
let admitted = self.sources.get(source)?;
Some(
admitted
.as_deref()
.is_some_and(|current| current.contains(normalised(text).as_str())),
)
}
pub(crate) fn refused_sources(&self) -> impl Iterator<Item = &str> {
self.sources
.iter()
.filter(|(_, admitted)| admitted.is_none())
.map(|(source, _)| source.as_str())
}
pub(crate) fn serves(&mut self, source: &str, text: &str) -> bool {
self.serves_with(source, text, None)
}
pub(crate) fn admits_unserved(&mut self, source: &str, text: &str) -> bool {
let admitted = self.decide(source, text);
if !admitted {
self.withheld += 1;
}
admitted
}
pub(crate) fn serves_with(
&mut self,
source: &str,
text: &str,
input_origin: Option<&crate::core::context_admission::origin::InputOrigin>,
) -> bool {
let served = self.decide(source, text);
if served {
Self::note_served(source, input_origin);
} else {
self.withheld += 1;
}
served
}
pub(crate) fn note_served(
source: &str,
input_origin: Option<&crate::core::context_admission::origin::InputOrigin>,
) {
if source.contains("://") || input_origin.is_some() {
crate::core::context_admission::capture::note_origin(
&crate::core::context_admission::stores::class_of_path(
Path::new(source),
input_origin,
),
);
}
}
#[cfg(test)]
pub(crate) fn withheld(&self) -> usize {
self.withheld
}
fn decide(&mut self, source: &str, text: &str) -> bool {
let external = source.contains("://");
let document = !external && crate::core::binary_detect::has_binary_extension(source);
let origins = Arc::clone(&self.origins);
let recorded = origins.get(source);
if let (true, Some(ChunkOrigin::Document(origin))) = (document, recorded) {
let label = self.root.join(source);
return self.authority_holds(source, |gate| {
origin.unchanged(&gate.root, source, gate.tool)
}) && self.admits_verbatim(text, label.to_str())
&& self.served_without_file_lineage();
}
let origin = match recorded {
Some(recorded) if external && self.protected() => match recorded {
ChunkOrigin::Provider(origin) => Some(&**origin),
ChunkOrigin::Delivery(origin) => {
if !origin.binds(text) {
return false;
}
if !self.authorities.contains_key(source) {
if self.lineage_checks >= MAX_LINEAGE_CHECKS {
return false;
}
self.lineage_checks += 1;
}
let served = self.authority_holds(source, |_| origin.still_admitted())
&& self.admits_verbatim(text, None);
if served {
origin.publish();
}
return served;
}
ChunkOrigin::Downstream(_) => return false,
ChunkOrigin::Document(_) => None,
},
_ => None,
};
if origin.is_none() && (external || document) {
return (external || self.root.join(source).is_file()) && self.originless(text);
}
if let Some(served) = self.decided(source, text) {
return served && (origin.is_none() || self.served_without_file_lineage());
}
if self.over_bound() {
return false;
}
let admitted = match origin {
Some(origin) => {
let deadline = *self
.provider_deadline
.get_or_insert_with(|| std::time::Instant::now() + PROVIDER_BUDGET);
if self.provider_checks >= MAX_PROVIDER_CHECKS
|| std::time::Instant::now() >= deadline
{
return false;
}
self.provider_checks += 1;
let admitted = self.read_provider(origin, deadline);
if admitted.is_none() && std::time::Instant::now() >= deadline {
return false;
}
admitted
}
None => self.read(source),
};
self.sources.insert(source.to_owned(), admitted);
self.decided(source, text).unwrap_or(false)
&& (origin.is_none() || self.served_without_file_lineage())
}
fn read_provider(
&self,
origin: &ProviderOrigin,
deadline: std::time::Instant,
) -> Option<String> {
let root = self.root.to_str()?;
let view = match crate::core::policy::runtime::with_project_source_view(root, || {
with_reuse_deadline_until(Some(deadline), || origin.current_chunks(root))
}) {
Ok(Ok(view)) => view,
Ok(Err(refusal)) => {
crate::core::context_admission::record_recheck_refusal(
&self.admission,
refusal.validity,
);
return None;
}
Err(_) => return None,
};
let mut current = String::new();
for chunk in &view.chunks {
let admitted = self.admission.admit(
&chunk.content,
Path::new(&chunk.file_path),
chunk.input_origin.as_ref(),
)?;
current.push_str(&normalised(&admitted));
current.push('\n');
}
Some(current)
}
fn read(&self, source: &str) -> Option<String> {
let path = self.root.join(source);
let read = crate::core::admitted_sources::read_with_admission(
self.root.to_str()?,
path.to_str()?,
self.tool,
&self.admission,
);
read.storable().map(normalised)
}
}
fn normalised(text: &str) -> String {
text.lines().collect::<Vec<_>>().join("\n")
}
pub fn chunk_served(gate: &mut Option<SourceGate>, chunk: &CodeChunk) -> bool {
gate.as_mut().is_none_or(|gate| {
gate.serves_with(
&chunk.file_path,
&chunk.content,
chunk.input_origin.as_ref(),
)
})
}
impl BM25Index {
#[must_use]
pub(crate) fn rooted_at(mut self, root: &Path) -> Self {
self.source_root = Some(root.to_path_buf());
self
}
pub fn source_gate(&self) -> Option<SourceGate> {
self.source_root
.as_deref()
.map(|root| SourceGate::new(root, SEARCH_TOOL).with_origins(&self.provider_origins))
}
pub fn search(&self, query: &str, top_k: usize) -> Vec<SearchResult> {
pinned(self.source_root.as_deref(), || {
self.search_gated(query, top_k, &mut self.source_gate())
})
.unwrap_or_default()
}
}
pub fn pinned<T>(root: Option<&Path>, search: impl FnOnce() -> T) -> Option<T> {
let Some(root) = root else {
return Some(search());
};
if !crate::core::workspace_authority::root_authorized(root) {
return None;
}
crate::core::policy::runtime::with_project_source_view(root.to_str()?, search).ok()
}
#[cfg(feature = "embeddings")]
pub(crate) fn overfetch(top_k: usize) -> usize {
top_k.saturating_add(top_k.clamp(16, MAX_SOURCES))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::core::policy::runtime::TestPolicyOverride;
#[test]
fn a_loaded_index_serves_a_chunk_only_while_its_source_admits_it() {
let _data = crate::core::data_dir::isolated_data_dir();
let _community = TestPolicyOverride::set(None);
let project = tempfile::tempdir().expect("project");
let root = project.path();
let withdrawn = root.join("billing.rs");
std::fs::write(
&withdrawn,
"pub fn settle_canary() -> u32 {\n let reconciliation = 7319;\n reconciliation\n}\n",
)
.expect("source");
std::fs::write(
root.join("ledger.rs"),
"pub fn settle_ledger() -> u32 {\n let reconciliation = 42;\n reconciliation\n}\n",
)
.expect("source");
let built = BM25Index::build_from_directory(root);
built.save(root).expect("persisted");
let files = |index: &BM25Index| -> Vec<String> {
index
.search("reconciliation", 10)
.into_iter()
.map(|hit| hit.file_path)
.collect()
};
let loaded = BM25Index::load(root).expect("index of the current policy");
assert!(loaded.source_root.is_some());
let served = files(&loaded);
assert!(served.iter().any(|file| file == "billing.rs"), "{served:?}");
assert!(served.iter().any(|file| file == "ledger.rs"), "{served:?}");
assert!(!loaded.search("settle_canary", 5).is_empty());
let unrooted = BM25Index::load_persisted(root).expect("persisted file");
std::fs::write(&withdrawn, "pub fn unrelated() {}\n").expect("rewrite");
let mut gate = loaded.source_gate();
let hits = loaded.search_gated("reconciliation", 10, &mut gate);
assert!(hits.iter().all(|hit| hit.file_path == "ledger.rs"));
assert!(!hits.is_empty());
assert!(gate.as_ref().is_some_and(|gate| gate.withheld() >= 1));
std::fs::remove_file(&withdrawn).expect("delete");
let served = files(&loaded);
assert!(!served.is_empty());
assert!(served.iter().all(|file| file == "ledger.rs"), "{served:?}");
assert!(loaded.search("settle_canary", 5).is_empty());
let mut full = SourceGate::new(root, SEARCH_TOOL);
for n in 0..MAX_SOURCES {
full.sources.insert(format!("known-{n}.rs"), None);
}
assert!(!full.serves("ledger.rs", "reconciliation"));
assert_eq!(full.sources.len(), MAX_SOURCES);
assert!(unrooted.source_root.is_none());
assert!(files(&unrooted).iter().any(|file| file == "billing.rs"));
}
}