use std::fs;
use std::path::{Path, PathBuf};
use crate::context::model::ContextFragment;
use crate::error::MemoryError;
use crate::limits::{
MAX_INGEST_FILES, MAX_INGEST_FILE_BYTES, MAX_TOTAL_INGEST_BYTES, MAX_TRANSCRIPT_BYTES,
};
#[derive(Debug, Clone, Default)]
pub struct IngestRoots {
roots: Vec<PathBuf>,
}
impl IngestRoots {
pub fn parse(value: &str) -> Result<Self, String> {
let mut roots = Vec::new();
for raw in std::env::split_paths(value) {
if raw.as_os_str().is_empty() {
continue;
}
let canonical = fs::canonicalize(&raw).map_err(|err| {
format!(
"VELESDB_MEMORY_INGEST_ROOTS entry '{}' could not be resolved: {err}",
raw.display()
)
})?;
if !canonical.is_dir() {
return Err(format!(
"VELESDB_MEMORY_INGEST_ROOTS entry '{}' is not a directory",
raw.display()
));
}
roots.push(canonical);
}
Ok(Self { roots })
}
#[must_use]
pub fn is_enabled(&self) -> bool {
!self.roots.is_empty()
}
fn contains(&self, candidate: &Path) -> bool {
self.roots.iter().any(|root| candidate.starts_with(root))
}
}
pub fn resolve_fragments(
fragments: &mut [ContextFragment],
roots: Option<&IngestRoots>,
) -> Result<(), MemoryError> {
for fragment in fragments.iter() {
if fragment.path.is_some() && (!fragment.content.is_empty() || fragment.media.is_some()) {
return Err(MemoryError::IngestPath(
"a fragment may set `path`, or `content`, or `media` — never more than one"
.to_owned(),
));
}
}
let indices: Vec<usize> = fragments
.iter()
.enumerate()
.filter(|(_, fragment)| fragment.path.is_some())
.map(|(index, _)| index)
.collect();
if indices.is_empty() {
return Ok(());
}
let Some(roots) = roots.filter(|roots| roots.is_enabled()) else {
return Err(MemoryError::IngestDisabled);
};
if indices.len() > MAX_INGEST_FILES {
return Err(MemoryError::ContextOverLimit(format!(
"request references {} files via `path`, exceeding the cap of {MAX_INGEST_FILES}",
indices.len()
)));
}
let mut total_bytes: usize = 0;
for index in indices {
let Some(requested) = fragments[index].path.take() else {
continue;
};
let content = resolve_one(&requested, roots, &mut total_bytes, MAX_INGEST_FILE_BYTES)?;
fragments[index].content = content;
}
Ok(())
}
pub fn resolve_transcript_path(
requested: &str,
roots: &IngestRoots,
) -> Result<String, MemoryError> {
let mut total_bytes: usize = 0;
resolve_one(requested, roots, &mut total_bytes, MAX_TRANSCRIPT_BYTES)
}
fn resolve_one(
requested: &str,
roots: &IngestRoots,
total_bytes: &mut usize,
file_cap: usize,
) -> Result<String, MemoryError> {
let requested_path = Path::new(requested);
if requested_path.is_relative() {
return Err(MemoryError::IngestPath(format!(
"path '{requested}' is relative; only absolute paths are accepted"
)));
}
let canonical = fs::canonicalize(requested_path).map_err(|err| {
MemoryError::IngestPath(format!("cannot resolve path '{requested}': {err}"))
})?;
if !roots.contains(&canonical) {
return Err(MemoryError::IngestOutsideRoots(requested.to_owned()));
}
let file_metadata = fs::metadata(&canonical)
.map_err(|err| MemoryError::IngestPath(format!("cannot stat path '{requested}': {err}")))?;
if !file_metadata.is_file() {
return Err(MemoryError::IngestPath(format!(
"path '{requested}' is not a regular file"
)));
}
let declared_len = usize::try_from(file_metadata.len()).unwrap_or(usize::MAX);
if declared_len > file_cap {
return Err(MemoryError::ContextOverLimit(format!(
"file '{requested}' is {declared_len} bytes, exceeding the cap of {file_cap} bytes"
)));
}
let running_total = total_bytes.saturating_add(declared_len);
if running_total > MAX_TOTAL_INGEST_BYTES {
return Err(MemoryError::ContextOverLimit(format!(
"ingesting '{requested}' would bring the request total to {running_total} bytes, \
exceeding the cap of {MAX_TOTAL_INGEST_BYTES} bytes"
)));
}
let bytes = fs::read(&canonical)
.map_err(|err| MemoryError::IngestPath(format!("cannot read path '{requested}': {err}")))?;
if bytes.len() > file_cap {
return Err(MemoryError::ContextOverLimit(format!(
"file '{requested}' grew to {} bytes while being read, exceeding the cap of \
{file_cap} bytes",
bytes.len()
)));
}
*total_bytes = total_bytes.saturating_add(bytes.len());
String::from_utf8(bytes).map_err(|err| {
let hint = magic_bytes_hint(err.as_bytes());
MemoryError::IngestPath(format!("path '{requested}' is not valid UTF-8{hint}"))
})
}
fn magic_bytes_hint(bytes: &[u8]) -> &'static str {
const PNG_MAGIC: [u8; 8] = [0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A];
const JPEG_MAGIC: [u8; 3] = [0xFF, 0xD8, 0xFF];
if bytes.starts_with(&PNG_MAGIC) || bytes.starts_with(&JPEG_MAGIC) {
" (looks like an image — use a media fragment instead of `path`)"
} else {
""
}
}
#[cfg(test)]
#[path = "ingest_tests.rs"]
mod tests;