use std::collections::{BTreeMap, HashMap, HashSet};
use std::path::Path;
use std::sync::Arc;
use serde::{Deserialize, Serialize};
use super::BM25Index;
use super::external_origin::ChunkOrigin;
use crate::core::content_chunk::ContentChunk;
const FILE: &str = "bm25_provider_origins.json";
const ARTIFACT_FILE: &str = "bm25_artifact_origins.json";
const VERSION: u8 = 2;
const MAX_ORIGINS: usize = 4096;
const MAX_ORIGIN_BYTES: usize = 32 * 1024;
const MAX_FILE_BYTES: usize = 16 * 1024 * 1024;
#[derive(Serialize, Deserialize)]
struct Recorded {
version: u8,
origins: BTreeMap<String, ChunkOrigin>,
}
#[derive(Debug, Clone, Default)]
pub struct ChunkOrigins(HashMap<String, ChunkOrigin>);
impl ChunkOrigins {
#[must_use]
pub fn len(&self) -> usize {
self.0.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.0.is_empty()
}
pub(crate) fn get(&self, source: &str) -> Option<&ChunkOrigin> {
self.0.get(source)
}
}
impl BM25Index {
pub fn ingest_content_chunks(
&mut self,
chunks: impl IntoIterator<Item = ContentChunk>,
) -> usize {
let mut count = 0usize;
for mut chunk in chunks {
if let Some(origin) = chunk.origin.take() {
self.record_origin(&chunk.file_path, ChunkOrigin::Provider(Box::new(origin)));
}
self.add_chunk(chunk.into());
count += 1;
}
if count > 0 {
self.finalize();
}
count
}
pub fn external_chunk_count(&self) -> usize {
self.chunks
.iter()
.filter(|c| c.file_path.contains("://"))
.count()
}
pub(crate) fn record_origin(&mut self, source: &str, origin: ChunkOrigin) {
let bounded = origin.fits(source)
&& (self.provider_origins.len() < MAX_ORIGINS
|| self.provider_origins.get(source).is_some())
&& serde_json::to_vec(&origin).is_ok_and(|bytes| bytes.len() <= MAX_ORIGIN_BYTES);
let origins = &mut Arc::make_mut(&mut self.provider_origins).0;
if bounded {
origins.insert(source.to_owned(), origin);
} else {
origins.remove(source);
}
}
pub(crate) fn bind_external_origins(&mut self, origins: Vec<(String, ChunkOrigin)>) {
for (source, origin) in origins {
let mut held = self.chunks.iter().filter(|chunk| chunk.file_path == source);
let bound = match (held.next(), held.next()) {
(Some(chunk), None) => Some(origin.bound_to(&chunk.content)),
_ => None,
};
match bound {
Some(origin) => self.record_origin(&source, origin),
None => {
Arc::make_mut(&mut self.provider_origins).0.remove(&source);
}
}
}
}
#[must_use]
pub(super) fn with_provider_origins(self, dir: &Path) -> Self {
self.with_origins_at(&dir.join(FILE))
}
pub(super) fn save_provider_origins(&self, dir: &Path) {
self.save_origins_at(dir, FILE);
}
#[must_use]
pub(crate) fn with_artifact_origins(self, index_file: &Path) -> Self {
self.with_origins_at(&index_file.with_file_name(ARTIFACT_FILE))
}
pub(crate) fn save_artifact_origins(&self, index_file: &Path) {
if let Some(dir) = index_file.parent() {
self.save_origins_at(dir, ARTIFACT_FILE);
}
}
fn with_origins_at(mut self, file: &Path) -> Self {
let held: HashSet<&str> = self
.chunks
.iter()
.map(|chunk| chunk.file_path.as_str())
.collect();
let origins: HashMap<String, ChunkOrigin> = read(file)
.unwrap_or_default()
.into_iter()
.filter(|(source, origin)| held.contains(source.as_str()) && origin.fits(source))
.collect();
self.provider_origins = Arc::new(ChunkOrigins(origins));
self
}
fn save_origins_at(&self, dir: &Path, name: &str) {
let path = dir.join(name);
let held: HashSet<&str> = if self.provider_origins.is_empty() {
HashSet::new()
} else {
self.chunks
.iter()
.map(|chunk| chunk.file_path.as_str())
.collect()
};
let origins: BTreeMap<String, ChunkOrigin> = self
.provider_origins
.0
.iter()
.filter(|(source, _)| held.contains(source.as_str()))
.map(|(source, origin)| (source.clone(), origin.clone()))
.collect();
let written = !origins.is_empty()
&& serde_json::to_vec(&Recorded {
version: VERSION,
origins,
})
.ok()
.filter(|bytes| bytes.len() <= MAX_FILE_BYTES)
.is_some_and(|bytes| {
let tmp = dir.join(format!("{name}.tmp"));
std::fs::write(&tmp, bytes).is_ok() && std::fs::rename(&tmp, &path).is_ok()
});
if !written {
let _ = std::fs::remove_file(&path);
}
}
}
fn read(path: &Path) -> Option<BTreeMap<String, ChunkOrigin>> {
if std::fs::metadata(path).ok()?.len() > MAX_FILE_BYTES as u64 {
return None;
}
let recorded: Recorded = serde_json::from_slice(&std::fs::read(path).ok()?).ok()?;
(recorded.version == VERSION && recorded.origins.len() <= MAX_ORIGINS)
.then_some(recorded.origins)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::core::providers::config_provider::{ConfigProvider, schema::ProviderConfig};
use crate::core::providers::hardened_http::redirect_tests::Server;
use crate::core::providers::provenance::BoundResult;
use crate::core::providers::provider_trait::ProviderParams;
#[test]
fn a_provider_chunk_is_served_only_while_its_origin_is_still_admitted() {
let _data = crate::core::data_dir::isolated_data_dir();
let _live_policy = crate::core::policy::runtime::lock_test_override();
let directory = tempfile::tempdir().expect("project");
let root = directory.path().canonicalize().expect("canonical root");
crate::core::workspace_trust::trust(&root).unwrap();
let configs = root.join(".lean-ctx/providers");
std::fs::create_dir_all(&configs).expect("provider directory");
let policy = root.join(".lean-ctx/policy.toml");
let pack = "name='index-origin'\nversion='1.0.0'\ndescription='fixture'\n";
std::fs::write(&policy, pack).expect("policy pack");
let body =
r#"[{"id":"1","title":"Orchidquartz outage","body":"Orchidquartz login path fails"}]"#;
let server = Server::new(200, None, body);
let config_text = format!(
"id='index-origin-fixture'\nname='Fixture'\nbase_url='{}'\n[auth]\ntype='none'\n\
[resources.issues]\nmethod='GET'\npath='/items'\n\
[resources.issues.response.mapping]\nid='id'\ntitle='title'\nbody='body'\n",
server.url
);
std::fs::write(configs.join("fixture.toml"), &config_text).expect("provider config");
let config: ProviderConfig = toml::from_str(&config_text).expect("config");
let provider = ConfigProvider::from_config(config).expect("provider");
let acquired = BoundResult::execute(&provider, "issues", &ProviderParams::default())
.expect("acquired");
let chunks = acquired.chunks(&acquired.result);
assert!(chunks.iter().all(|chunk| chunk.origin.is_some()));
let mut written = BM25Index::new();
written.admission_policy = Some(
crate::core::context_admission::stores::StoreAdmission::current()
.digest()
.to_owned(),
);
assert_eq!(written.ingest_content_chunks(chunks), 1);
written.save(&root).expect("persisted");
let loaded = BM25Index::load(&root).expect("index of the current policy");
assert_eq!(loaded.provider_origins.len(), 1);
let served = || !loaded.search("orchidquartz", 5).is_empty();
assert!(served(), "admitted origin: served under the pack");
std::fs::write(
&policy,
format!("{pack}[context]\ndeny_tools = ['ctx_provider']\n"),
)
.expect("tightened pack");
assert!(!served(), "the pack denies the provider");
std::fs::write(&policy, pack).expect("pack restored");
assert!(served());
server.respond(403, "denied");
assert!(!served(), "the provider no longer returns the item");
std::fs::remove_file(&policy).expect("pack removed");
let asked = server.count();
assert!(served());
assert_eq!(server.count(), asked, "no provider round trip");
}
}