use crate::{Edge, EdgeKind, FactSet, Node, NodeKind, Provenance, Span};
pub(crate) const EXTRACT_VERSION: u32 = 11
+ if cfg!(feature = "pdf-text") { 100 } else { 0 }
+ if cfg!(feature = "image-ocr") { 200 } else { 0 }
+ if cfg!(feature = "audio-metadata") {
400
} else {
0
};
const MAX_CONTENT: usize = 1500;
#[cfg(feature = "pdf-text")]
const MAX_PDF_BYTES: usize = 20 * 1024 * 1024;
#[cfg(any(feature = "image-ocr", feature = "image-vision"))]
const MAX_IMAGE_PIXELS: u64 = 4096 * 4096;
pub trait Extractor {
fn extract(&self, path: &str, blob_id: &str, bytes: &[u8]) -> FactSet;
fn env_tag(&self) -> u64 {
media_env_tag()
}
}
#[allow(clippy::struct_excessive_bools)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct IngestConfig {
pub prose: bool,
pub pdf: bool,
pub ocr: bool,
pub vision: bool,
pub audio: bool,
}
impl Default for IngestConfig {
fn default() -> Self {
Self {
prose: true,
pdf: true,
ocr: true,
vision: true,
audio: true,
}
}
}
impl IngestConfig {
fn disabled_bits(self) -> u64 {
u64::from(!self.prose) | (u64::from(!self.pdf) << 1) | (u64::from(!self.ocr) << 2)
}
#[must_use]
pub fn generates(self, kind: crate::media::MediaKind) -> bool {
match kind {
crate::media::MediaKind::Audio => self.audio,
crate::media::MediaKind::Vision => self.vision,
}
}
}
#[derive(Debug, Clone, Copy, Default)]
pub struct Registry {
pub ingest: IngestConfig,
}
impl Registry {
#[must_use]
pub fn new(ingest: IngestConfig) -> Self {
Self { ingest }
}
}
impl Extractor for Registry {
fn extract(&self, path: &str, blob_id: &str, bytes: &[u8]) -> FactSet {
let mut facts = extract_facts(path, blob_id, bytes, self.ingest);
crate::markers::augment(&mut facts, path, blob_id, bytes);
facts
}
fn env_tag(&self) -> u64 {
let media = media_env_tag();
let disabled = self.ingest.disabled_bits();
if disabled == 0 {
media
} else {
let mut h = 0xcbf2_9ce4_8422_2325u64;
for b in media
.to_le_bytes()
.into_iter()
.chain(disabled.to_le_bytes())
{
h ^= u64::from(b);
h = h.wrapping_mul(0x0000_0100_0000_01b3);
}
h
}
}
}
fn extract_facts(path: &str, blob_id: &str, bytes: &[u8], ingest: IngestConfig) -> FactSet {
if crate::config_keys::is_config_path(path) {
return config_facts(path, blob_id, bytes, ingest);
}
if is_dockerfile(path) {
return dockerfile_facts(path, blob_id, bytes, ingest);
}
if crate::media::is_audio(path) {
return audio_facts(path, blob_id, bytes, ingest);
}
let ext = extension(path);
match ext.as_deref() {
Some("rs") => rust_facts(path, blob_id, bytes, ingest),
Some(ext) => tag_facts(path, blob_id, bytes, ext, ingest).unwrap_or_else(|| {
FactSet::new().with_node(file_node(path, blob_id, bytes, None, ingest))
}),
None => FactSet::new().with_node(file_node(path, blob_id, bytes, None, ingest)),
}
}
pub(crate) fn extension(path: &str) -> Option<String> {
let name = path.rsplit('/').next().unwrap_or(path);
name.rsplit_once('.')
.map(|(_, ext)| ext.to_ascii_lowercase())
}
fn file_key(path: &str) -> String {
format!("file:{path}")
}
fn file_node(
path: &str,
blob_id: &str,
bytes: &[u8],
lang: Option<&str>,
ingest: IngestConfig,
) -> Node {
let name = path.rsplit('/').next().unwrap_or(path).to_owned();
let lines = bytes
.iter()
.fold(0usize, |n, &b| n + usize::from(b == b'\n'));
let end = u32::try_from(bytes.len()).unwrap_or(u32::MAX);
let mut meta = serde_json::json!({ "bytes": bytes.len(), "lines": lines });
let content = if ingest.prose && is_prose(path) {
cap_content(&String::from_utf8_lossy(bytes))
} else if let Some(text) = ingest.pdf.then(|| pdf_content(path, bytes)).flatten() {
cap_content(&text)
} else if let Some(text) = image_content(path, bytes, ingest) {
cap_content(&text)
} else {
String::new()
};
if !content.is_empty() {
meta["content"] = serde_json::Value::from(content);
}
Node {
key: file_key(path),
kind: NodeKind::File,
name,
path: Some(path.to_owned()),
lang: lang.map(ToOwned::to_owned),
blob_hash: Some(blob_id.to_owned()),
span: Some(Span::new(0, end)),
provenance: Provenance::Derived,
meta,
}
}
fn config_facts(path: &str, blob_id: &str, bytes: &[u8], ingest: IngestConfig) -> FactSet {
let mut facts = FactSet::new().with_node(file_node(path, blob_id, bytes, None, ingest));
let file = file_key(path);
let mut by_key: std::collections::BTreeMap<String, String> = std::collections::BTreeMap::new();
for ck in crate::config_keys::flatten(path, bytes) {
by_key.insert(ck.key, ck.value);
}
for (key, value) in by_key {
let node_key = format!("cfgkey:{path}#{key}");
let value = if crate::config_keys::is_secret_key(&key) {
crate::config_keys::REDACTED.to_owned()
} else {
value
};
let mut node = Node::new(
node_key.clone(),
NodeKind::Other(crate::config_keys::KIND.into()),
key.clone(),
);
node.path = Some(path.to_owned());
node.blob_hash = Some(blob_id.to_owned());
node.meta = serde_json::json!({ "key": key, "value": value });
facts = facts.with_node(node).with_edge(Edge::derived(
file.clone(),
node_key,
EdgeKind::Contains,
));
}
facts
}
pub(crate) const IMAGE_REF_KIND: &str = "image_ref";
fn is_dockerfile(path: &str) -> bool {
let base = path.rsplit('/').next().unwrap_or(path).to_ascii_lowercase();
base == "dockerfile"
|| base == "containerfile"
|| base.starts_with("dockerfile.")
|| base.ends_with(".dockerfile")
}
fn dockerfile_facts(path: &str, blob_id: &str, bytes: &[u8], ingest: IngestConfig) -> FactSet {
let mut facts = FactSet::new().with_node(file_node(path, blob_id, bytes, None, ingest));
let file = file_key(path);
let text = String::from_utf8_lossy(bytes);
let mut stages: std::collections::HashSet<String> = std::collections::HashSet::new();
let mut idx = 0usize;
for line in text.lines() {
let Some(rest) = strip_from_prefix(line.trim()) else {
continue;
};
let (image, stage) = parse_from(rest);
let is_internal_stage = stages.contains(&image.to_ascii_lowercase());
if let Some(s) = stage {
stages.insert(s.to_ascii_lowercase());
}
if image.is_empty() || image.eq_ignore_ascii_case("scratch") || is_internal_stage {
continue;
}
let (name, tag, digest) = split_image(image);
let node_key = format!("imageref:{path}#{idx}");
idx += 1;
let mut node = Node::new(
node_key.clone(),
NodeKind::Other(IMAGE_REF_KIND.into()),
image.to_owned(),
);
node.path = Some(path.to_owned());
node.blob_hash = Some(blob_id.to_owned());
node.meta = serde_json::json!({ "image": name, "tag": tag, "digest": digest });
facts = facts.with_node(node).with_edge(Edge::derived(
file.clone(),
node_key,
EdgeKind::References,
));
}
facts
}
fn audio_facts(path: &str, blob_id: &str, bytes: &[u8], ingest: IngestConfig) -> FactSet {
let facts = FactSet::new().with_node(file_node(path, blob_id, bytes, None, ingest));
let Some(node) = audio_stream_node(path, blob_id, bytes) else {
return facts;
};
let node_key = node.key.clone();
facts
.with_node(node)
.with_edge(Edge::derived(file_key(path), node_key, EdgeKind::Contains))
}
#[cfg(feature = "audio-metadata")]
fn audio_stream_node(path: &str, blob_id: &str, bytes: &[u8]) -> Option<Node> {
let facts = crate::audio::read(bytes, extension(path).as_deref())?;
let mut meta = serde_json::to_value(&facts).ok()?;
meta["content"] = serde_json::Value::from(cap_content(&facts.summary()));
let name = path.rsplit('/').next().unwrap_or(path).to_owned();
let mut node = Node::new(
format!("audio:{path}"),
NodeKind::Other(crate::audio::AUDIO_STREAM_KIND.into()),
name,
);
node.path = Some(path.to_owned());
node.blob_hash = Some(blob_id.to_owned());
node.span = Some(Span::new(0, u32::try_from(bytes.len()).unwrap_or(u32::MAX)));
node.meta = meta;
Some(node)
}
#[cfg(not(feature = "audio-metadata"))]
fn audio_stream_node(_path: &str, _blob_id: &str, _bytes: &[u8]) -> Option<Node> {
None
}
fn strip_from_prefix(line: &str) -> Option<&str> {
let b = line.as_bytes();
(b.len() >= 5 && b[..4].eq_ignore_ascii_case(b"from") && b[4].is_ascii_whitespace())
.then(|| line[5..].trim_start())
}
fn parse_from(rest: &str) -> (&str, Option<&str>) {
let image = rest
.split_whitespace()
.find(|t| !t.starts_with("--"))
.unwrap_or("");
let mut toks = rest.split_whitespace();
let mut stage = None;
while let Some(t) = toks.next() {
if t.eq_ignore_ascii_case("as") {
stage = toks.next();
break;
}
}
(image, stage)
}
fn split_image(image: &str) -> (String, Option<String>, Option<String>) {
if let Some((name, digest)) = image.split_once('@') {
return (name.to_owned(), None, Some(digest.to_owned()));
}
let seg = image.rfind('/').map_or(0, |i| i + 1);
if let Some(colon) = image[seg..].find(':') {
let at = seg + colon;
return (
image[..at].to_owned(),
Some(image[at + 1..].to_owned()),
None,
);
}
(image.to_owned(), None, None)
}
fn doc_comment_body(raw: &str) -> Option<String> {
let t = raw.trim();
if t.starts_with("//!") || (t.starts_with("///") && !t.starts_with("////")) {
return Some(t[3..].trim().to_owned());
}
if (t.starts_with("/**") || t.starts_with("/*!")) && t.ends_with("*/") {
let end = t.len() - 2;
let inner = if end >= 3 { &t[3..end] } else { "" };
let cleaned: Vec<&str> = inner
.lines()
.map(|l| l.trim().trim_start_matches('*').trim())
.filter(|l| !l.is_empty())
.collect();
return Some(cleaned.join(" "));
}
None
}
#[cfg(feature = "pdf-text")]
fn pdf_content(path: &str, bytes: &[u8]) -> Option<String> {
if extension(path).as_deref() != Some("pdf") || bytes.len() > MAX_PDF_BYTES {
return None;
}
let owned = bytes.to_vec();
let text = std::panic::catch_unwind(move || pdf_extract::extract_text_from_mem(&owned).ok())
.ok()
.flatten()?;
(!text.trim().is_empty()).then_some(text)
}
#[cfg(not(feature = "pdf-text"))]
fn pdf_content(_path: &str, _bytes: &[u8]) -> Option<String> {
None
}
#[cfg(feature = "image-ocr")]
fn image_content(path: &str, bytes: &[u8], ingest: IngestConfig) -> Option<String> {
if !ingest.ocr || !crate::media::is_image(path) || bytes.len() > crate::media::MAX_IMAGE_BYTES {
return None;
}
ocr_content(bytes)
}
#[cfg(not(feature = "image-ocr"))]
fn image_content(_path: &str, _bytes: &[u8], _ingest: IngestConfig) -> Option<String> {
None
}
#[cfg(any(feature = "image-ocr", feature = "image-vision"))]
pub(crate) fn image_dimensions_ok(bytes: &[u8]) -> bool {
let Ok(reader) = image::ImageReader::new(std::io::Cursor::new(bytes)).with_guessed_format()
else {
return false;
};
match reader.into_dimensions() {
Ok((w, h)) => u64::from(w) * u64::from(h) <= MAX_IMAGE_PIXELS,
Err(_) => false,
}
}
#[cfg(feature = "image-ocr")]
fn ocr_content(bytes: &[u8]) -> Option<String> {
let model = crate::model_choice::resolve(crate::model_choice::ModelTask::Ocr)
.ok()?
.model?;
let dir = crate::models::model_dir(model);
let detection = dir.join("text-detection.rten");
let recognition = dir.join("text-recognition.rten");
if !detection.exists() || !recognition.exists() || !image_dimensions_ok(bytes) {
return None;
}
let text = std::panic::catch_unwind(|| run_ocr(&detection, &recognition, bytes))
.ok()
.flatten()?;
(!text.trim().is_empty()).then_some(text)
}
#[cfg(feature = "image-ocr")]
fn run_ocr(
detection: &std::path::Path,
recognition: &std::path::Path,
bytes: &[u8],
) -> Option<String> {
use ocrs::{ImageSource, OcrEngine, OcrEngineParams};
let detection_model = rten::Model::load_file(detection).ok()?;
let recognition_model = rten::Model::load_file(recognition).ok()?;
let engine = OcrEngine::new(OcrEngineParams {
detection_model: Some(detection_model),
recognition_model: Some(recognition_model),
..Default::default()
})
.ok()?;
let img = image::load_from_memory(bytes).ok()?.into_rgb8();
let source = ImageSource::from_bytes(img.as_raw(), img.dimensions()).ok()?;
let input = engine.prepare_input(source).ok()?;
engine.get_text(&input).ok()
}
#[must_use]
pub fn release_media_engines() -> bool {
let vision = release_vlm_engine();
let audio = release_asr_engine();
let backend = release_llama_backend();
vision || audio || backend
}
#[cfg(any(feature = "image-vision", feature = "audio-transcribe"))]
fn release_llama_backend() -> bool {
rto_llama::backend::release_shared_backend()
}
#[cfg(not(any(feature = "image-vision", feature = "audio-transcribe")))]
fn release_llama_backend() -> bool {
false
}
fn release_vlm_engine() -> bool {
crate::media::producers::release_vlm_engine()
}
fn release_asr_engine() -> bool {
crate::media::producers::release_asr_engine()
}
#[derive(Debug)]
pub struct MediaEngineGuard {
_private: (),
}
impl MediaEngineGuard {
#[must_use]
pub const fn hold() -> Self {
Self { _private: () }
}
}
impl Drop for MediaEngineGuard {
fn drop(&mut self) {
let _released = release_media_engines();
}
}
#[cfg(all(feature = "image-ocr", not(feature = "image-vision")))]
fn vlm_content(_bytes: &[u8]) -> Option<String> {
None
}
#[cfg(feature = "image-ocr")]
pub(crate) fn media_env_tag() -> u64 {
let Some(model) = crate::model_choice::resolve(crate::model_choice::ModelTask::Ocr)
.ok()
.and_then(|choice| choice.model)
else {
return 0;
};
let mut hash: u64 = 0xcbf2_9ce4_8422_2325;
if fold_installed_model(&mut hash, model) {
hash | 1
} else {
0
}
}
#[cfg(feature = "image-ocr")]
fn fold_installed_model(hash: &mut u64, name: &str) -> bool {
let Some(variant) = crate::models::find(name)
.and_then(|spec| spec.variant_for(crate::models::Platform::host()))
else {
return false;
};
let dir = crate::models::model_dir(name);
if !variant.files.iter().all(|f| dir.join(f.name).exists()) {
return false;
}
for file in variant.files {
for b in file.sha256.bytes() {
*hash ^= u64::from(b);
*hash = hash.wrapping_mul(0x0000_0100_0000_01b3);
}
}
true
}
#[cfg(not(feature = "image-ocr"))]
pub(crate) fn media_env_tag() -> u64 {
0
}
fn is_prose(path: &str) -> bool {
matches!(
extension(path).as_deref(),
Some("md" | "markdown" | "txt" | "rst" | "adoc")
)
}
fn cap_content(text: &str) -> String {
let mut out = String::with_capacity(text.len().min(MAX_CONTENT));
let mut chars = 0usize;
let mut last_was_space = true;
for c in text.chars() {
if chars >= MAX_CONTENT {
break;
}
if c.is_whitespace() {
if !last_was_space {
out.push(' ');
chars += 1;
last_was_space = true;
}
} else {
out.push(c);
chars += 1;
last_was_space = false;
}
}
out.trim().to_owned()
}
#[derive(Debug, Clone, Copy, Default)]
pub struct FileNodeExtractor;
impl Extractor for FileNodeExtractor {
fn extract(&self, path: &str, blob_id: &str, bytes: &[u8]) -> FactSet {
FactSet::new().with_node(file_node(
path,
blob_id,
bytes,
None,
IngestConfig::default(),
))
}
}
#[derive(Debug, Clone, Copy, Default)]
pub struct RustExtractor;
impl Extractor for RustExtractor {
fn extract(&self, path: &str, blob_id: &str, bytes: &[u8]) -> FactSet {
rust_facts(path, blob_id, bytes, IngestConfig::default())
}
}
fn rust_facts(path: &str, blob_id: &str, bytes: &[u8], ingest: IngestConfig) -> FactSet {
let mut parser = tree_sitter::Parser::new();
if parser
.set_language(&tree_sitter_rust::LANGUAGE.into())
.is_err()
{
return FactSet::new().with_node(file_node(path, blob_id, bytes, None, ingest));
}
let Some(tree) = parser.parse(bytes, None) else {
return FactSet::new().with_node(file_node(path, blob_id, bytes, None, ingest));
};
let mut walk = RustWalk {
path,
blob_id,
src: bytes,
nodes: vec![file_node(path, blob_id, bytes, Some("rust"), ingest)],
edges: Vec::new(),
};
let root = tree.root_node();
let mut cursor = root.walk();
let children: Vec<_> = root.children(&mut cursor).collect();
for child in children {
walk.visit(child, &[]);
}
walk.synthesize_config_keys(root);
walk.nodes.sort_by(|a, b| a.key.cmp(&b.key));
walk.edges
.sort_by(|a, b| (a.kind.as_str(), &a.src, &a.dst).cmp(&(b.kind.as_str(), &b.src, &b.dst)));
FactSet {
nodes: walk.nodes,
edges: walk.edges,
}
}
struct Scope {
seg: String,
key: Option<String>,
}
struct FieldDef {
name: String,
type_idents: Vec<String>,
}
const TRANSPARENT_WRAPPERS: &[&str] = &[
"Option", "Box", "Arc", "Rc", "Cow", "RefCell", "Cell", "Mutex", "RwLock",
];
const COLLECTION_WRAPPERS: &[&str] = &[
"Vec", "VecDeque", "HashMap", "BTreeMap", "HashSet", "BTreeSet", "IndexMap",
];
fn core_type_name(type_idents: &[String]) -> Option<String> {
type_idents
.iter()
.find(|t| !TRANSPARENT_WRAPPERS.contains(&t.as_str()))
.or_else(|| type_idents.first())
.cloned()
}
fn recursion_target<'a>(
type_idents: &'a [String],
known: &std::collections::BTreeMap<String, StructDef>,
) -> Option<&'a str> {
for t in type_idents {
if COLLECTION_WRAPPERS.contains(&t.as_str()) {
return None;
}
if TRANSPARENT_WRAPPERS.contains(&t.as_str()) {
continue;
}
return known.contains_key(t).then_some(t.as_str());
}
None
}
struct StructDef {
fields: Vec<FieldDef>,
is_root: bool,
}
const MAX_CONFIG_DEPTH: usize = 16;
fn expand_config_keys(
table: &std::collections::BTreeMap<String, StructDef>,
struct_name: &str,
prefix: &str,
root: &str,
visited: &mut std::collections::BTreeSet<String>,
depth: usize,
out: &mut std::collections::BTreeMap<String, String>,
) {
let Some(def) = table.get(struct_name) else {
return;
};
for f in &def.fields {
let key = if prefix.is_empty() {
f.name.clone()
} else {
format!("{prefix}.{}", f.name)
};
match recursion_target(&f.type_idents, table) {
Some(inner) if depth < MAX_CONFIG_DEPTH && !visited.contains(inner) => {
visited.insert(inner.to_owned());
expand_config_keys(table, inner, &key, root, visited, depth + 1, out);
visited.remove(inner);
}
_ => {
out.entry(key).or_insert_with(|| root.to_owned());
}
}
}
}
struct RustWalk<'a> {
path: &'a str,
blob_id: &'a str,
src: &'a [u8],
nodes: Vec<Node>,
edges: Vec<Edge>,
}
impl RustWalk<'_> {
fn visit(&mut self, node: tree_sitter::Node, scope: &[Scope]) {
match node.kind() {
"function_item" => self.visit_symbol(node, scope, NodeKind::Fn, true),
"struct_item" | "union_item" => self.visit_symbol(node, scope, NodeKind::Struct, false),
"enum_item" => self.visit_symbol(node, scope, NodeKind::Enum, false),
"trait_item" => self.visit_symbol(node, scope, NodeKind::Trait, false),
"mod_item" => self.visit_symbol(node, scope, NodeKind::Module, false),
"type_item" => self.visit_symbol(node, scope, NodeKind::Other("type".into()), false),
"macro_definition" => {
self.visit_symbol(node, scope, NodeKind::Other("macro".into()), false);
}
"impl_item" => self.visit_impl(node, scope),
"use_declaration" => self.visit_use(node),
_ => self.visit_children(node, scope),
}
}
fn visit_children(&mut self, node: tree_sitter::Node, scope: &[Scope]) {
let mut cursor = node.walk();
let children: Vec<_> = node.named_children(&mut cursor).collect();
for child in children {
self.visit(child, scope);
}
}
fn visit_symbol(
&mut self,
node: tree_sitter::Node,
scope: &[Scope],
kind: NodeKind,
collect_calls: bool,
) {
let Some(name) = self.field_text(node, "name") else {
return self.visit_children(node, scope);
};
let qualified = qualify(scope, &name);
let key = format!("sym:rust:{}#{qualified}", self.path);
let mut meta = serde_json::Map::new();
if collect_calls {
let mut calls = Vec::new();
self.collect_calls(node, &mut calls);
calls.sort();
calls.dedup();
if !calls.is_empty() {
meta.insert("calls".into(), serde_json::Value::from(calls));
}
}
if let Some(doc) = self.doc_comment(node) {
meta.insert("content".into(), serde_json::Value::from(doc));
}
if matches!(node.kind(), "struct_item" | "union_item") {
let defs = self.struct_fields(node);
if !defs.is_empty() {
let names: Vec<&str> = defs.iter().map(|f| f.name.as_str()).collect();
meta.insert("fields".into(), serde_json::Value::from(names));
let types: serde_json::Map<String, serde_json::Value> = defs
.iter()
.filter_map(|f| {
core_type_name(&f.type_idents).map(|t| (f.name.clone(), t.into()))
})
.collect();
if !types.is_empty() {
meta.insert("field_types".into(), serde_json::Value::Object(types));
}
}
if self.has_config_marker(node) {
meta.insert("config_root".into(), serde_json::Value::Bool(true));
}
}
self.nodes.push(Node {
key: key.clone(),
kind,
name,
path: Some(self.path.to_owned()),
lang: Some("rust".to_owned()),
blob_hash: Some(self.blob_id.to_owned()),
span: Some(span(node)),
provenance: Provenance::Derived,
meta: serde_json::Value::Object(meta),
});
self.link_parent(&key, scope);
let child_scope = extend(scope, &self.simple(node, "name"), Some(key));
self.recurse_body(node, &child_scope);
}
fn doc_comment(&self, node: tree_sitter::Node) -> Option<String> {
let mut parts: Vec<String> = Vec::new();
let mut prev = node.prev_sibling();
while let Some(n) = prev {
match n.kind() {
"line_comment" | "block_comment" => match doc_comment_body(self.text(n)) {
Some(body) => {
parts.push(body);
prev = n.prev_sibling();
}
None => break,
},
"attribute_item" => prev = n.prev_sibling(),
_ => break,
}
}
if parts.is_empty() {
return None;
}
parts.reverse();
let joined = cap_content(&parts.join(" "));
(!joined.is_empty()).then_some(joined)
}
fn visit_impl(&mut self, node: tree_sitter::Node, scope: &[Scope]) {
let type_name = self
.field_text(node, "type")
.unwrap_or_else(|| "impl".to_owned());
let child_scope = extend(scope, &type_name, None);
self.recurse_body(node, &child_scope);
}
fn visit_use(&mut self, node: tree_sitter::Node) {
let Some(arg) = node.child_by_field_name("argument") else {
return;
};
let text: String = self
.text(arg)
.chars()
.filter(|c| !c.is_whitespace())
.collect();
if text.is_empty() {
return;
}
let key = format!("import:rust:{text}");
self.nodes.push(Node {
key: key.clone(),
kind: NodeKind::Other("import".into()),
name: text,
path: None,
lang: Some("rust".to_owned()),
blob_hash: None,
span: None,
provenance: Provenance::Derived,
meta: serde_json::Value::Null,
});
self.edges
.push(Edge::derived(file_key(self.path), key, EdgeKind::Imports));
}
fn link_parent(&mut self, key: &str, scope: &[Scope]) {
if let Some(parent) = scope.iter().rev().find_map(|s| s.key.as_deref()) {
self.edges.push(Edge::derived(
parent.to_owned(),
key.to_owned(),
EdgeKind::Contains,
));
} else {
self.edges.push(Edge::derived(
file_key(self.path),
key.to_owned(),
EdgeKind::Defines,
));
}
}
fn struct_fields(&self, node: tree_sitter::Node) -> Vec<FieldDef> {
let mut out = Vec::new();
let mut cursor = node.walk();
for child in node.named_children(&mut cursor) {
if child.kind() == "field_declaration_list" {
let mut inner = child.walk();
for field in child.named_children(&mut inner) {
if field.kind() == "field_declaration"
&& let Some(name) = field.child_by_field_name("name")
{
let type_idents = field
.child_by_field_name("type")
.map(|t| self.type_idents(t))
.unwrap_or_default();
out.push(FieldDef {
name: self.text(name).to_owned(),
type_idents,
});
}
}
}
}
out
}
fn type_idents(&self, ty: tree_sitter::Node) -> Vec<String> {
let mut out = Vec::new();
self.collect_type_idents(ty, &mut out);
out
}
fn collect_type_idents(&self, node: tree_sitter::Node, out: &mut Vec<String>) {
if matches!(node.kind(), "type_identifier" | "primitive_type") {
out.push(self.text(node).to_owned());
}
let mut cursor = node.walk();
for child in node.named_children(&mut cursor) {
self.collect_type_idents(child, out);
}
}
fn has_config_marker(&self, node: tree_sitter::Node) -> bool {
const MARKER: &str = "@rto:config";
let mut prev = node.prev_sibling();
while let Some(n) = prev {
match n.kind() {
"line_comment" | "block_comment" | "attribute_item" => {
if self.text(n).contains(MARKER) {
return true;
}
prev = n.prev_sibling();
}
_ => break,
}
}
false
}
fn recurse_body(&mut self, node: tree_sitter::Node, scope: &[Scope]) {
let mut cursor = node.walk();
let children: Vec<_> = node.named_children(&mut cursor).collect();
for child in children {
match child.kind() {
"declaration_list" | "field_declaration_list" | "trait_body" => {
self.visit_children(child, scope);
}
_ => {}
}
}
}
fn collect_calls(&self, node: tree_sitter::Node, out: &mut Vec<String>) {
let mut cursor = node.walk();
for child in node.named_children(&mut cursor) {
if child.kind() == "call_expression"
&& let Some(func) = child.child_by_field_name("function")
&& let Some(name) = self.callee_name(func)
{
out.push(name);
}
self.collect_calls(child, out);
}
}
fn callee_name(&self, func: tree_sitter::Node) -> Option<String> {
match func.kind() {
"identifier" => Some(self.text(func).to_owned()),
"scoped_identifier" => {
let name = func.child_by_field_name("name")?;
let qualifier = func
.child_by_field_name("path")
.and_then(|p| self.text(p).rsplit("::").next().map(str::to_owned));
Some(qualify_callee(qualifier.as_deref(), self.text(name)))
}
"field_expression" => {
let name = func.child_by_field_name("field")?;
let on_self = func
.child_by_field_name("value")
.is_some_and(|v| self.text(v) == "self");
Some(qualify_callee(on_self.then_some("Self"), self.text(name)))
}
_ => None,
}
}
fn synthesize_config_keys(&mut self, root: tree_sitter::Node) {
let table = self.collect_struct_defs(root);
let mut keys: std::collections::BTreeMap<String, String> =
std::collections::BTreeMap::new();
for (name, def) in &table {
if !def.is_root {
continue;
}
let mut visited = std::collections::BTreeSet::new();
visited.insert(name.clone());
expand_config_keys(&table, name, "", name, &mut visited, 0, &mut keys);
}
let file = file_key(self.path);
for (dotted, root_name) in keys {
let node_key = format!("cfgkey:{}#{dotted}", self.path);
let mut node = Node::new(
node_key.clone(),
NodeKind::Other(crate::config_keys::KIND.into()),
dotted.clone(),
);
node.path = Some(self.path.to_owned());
node.blob_hash = Some(self.blob_id.to_owned());
node.meta = serde_json::json!({
"key": dotted,
"source": "struct",
"struct": root_name,
});
self.edges.push(Edge::derived(
file.clone(),
node_key.clone(),
EdgeKind::Contains,
));
self.nodes.push(node);
}
}
fn collect_struct_defs(
&self,
root: tree_sitter::Node,
) -> std::collections::BTreeMap<String, StructDef> {
let mut out = std::collections::BTreeMap::new();
self.collect_struct_defs_into(root, &mut out);
out
}
fn collect_struct_defs_into(
&self,
node: tree_sitter::Node,
out: &mut std::collections::BTreeMap<String, StructDef>,
) {
if matches!(node.kind(), "struct_item" | "union_item")
&& let Some(name) = self.field_text(node, "name")
{
out.entry(name.clone()).or_insert_with(|| StructDef {
fields: self.struct_fields(node),
is_root: self.has_config_marker(node),
});
}
let mut cursor = node.walk();
for child in node.named_children(&mut cursor) {
self.collect_struct_defs_into(child, out);
}
}
fn text(&self, node: tree_sitter::Node) -> &str {
node.utf8_text(self.src).unwrap_or("")
}
fn field_text(&self, node: tree_sitter::Node, field: &str) -> Option<String> {
node.child_by_field_name(field)
.map(|n| self.text(n).to_owned())
}
fn simple(&self, node: tree_sitter::Node, field: &str) -> String {
self.field_text(node, field).unwrap_or_default()
}
}
struct TagLang {
lang: &'static str,
grammar_key: &'static str,
language: tree_sitter::Language,
query: std::borrow::Cow<'static, str>,
}
#[allow(clippy::too_many_lines)]
fn tag_lang_for(ext: &str) -> Option<TagLang> {
use std::borrow::Cow;
let ts_query = || -> Cow<'static, str> {
Cow::Owned(format!(
"{}\n{}",
tree_sitter_javascript::TAGS_QUERY,
tree_sitter_typescript::TAGS_QUERY
))
};
let borrowed = |q: &'static str| -> Cow<'static, str> { Cow::Borrowed(q) };
let (lang, language, query): (&str, tree_sitter::Language, Cow<'static, str>) = match ext {
"py" | "pyi" => (
"python",
tree_sitter_python::LANGUAGE.into(),
borrowed(tree_sitter_python::TAGS_QUERY),
),
"js" | "jsx" | "mjs" | "cjs" => (
"javascript",
tree_sitter_javascript::LANGUAGE.into(),
borrowed(tree_sitter_javascript::TAGS_QUERY),
),
"ts" | "mts" | "cts" => (
"typescript",
tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into(),
ts_query(),
),
"tsx" => (
"tsx",
tree_sitter_typescript::LANGUAGE_TSX.into(),
ts_query(),
),
"go" => (
"go",
tree_sitter_go::LANGUAGE.into(),
borrowed(tree_sitter_go::TAGS_QUERY),
),
"rb" => (
"ruby",
tree_sitter_ruby::LANGUAGE.into(),
borrowed(tree_sitter_ruby::TAGS_QUERY),
),
"java" => (
"java",
tree_sitter_java::LANGUAGE.into(),
borrowed(tree_sitter_java::TAGS_QUERY),
),
"c" | "h" => (
"c",
tree_sitter_c::LANGUAGE.into(),
borrowed(tree_sitter_c::TAGS_QUERY),
),
"cc" | "cpp" | "cxx" | "hpp" | "hh" | "hxx" => (
"cpp",
tree_sitter_cpp::LANGUAGE.into(),
borrowed(tree_sitter_cpp::TAGS_QUERY),
),
"cs" => (
"csharp",
tree_sitter_c_sharp::LANGUAGE.into(),
borrowed(include_str!("queries/csharp/tags.scm")),
),
"php" => (
"php",
tree_sitter_php::LANGUAGE_PHP.into(),
borrowed(tree_sitter_php::TAGS_QUERY),
),
"scala" | "sc" => (
"scala",
tree_sitter_scala::LANGUAGE.into(),
borrowed(include_str!("queries/scala/tags.scm")),
),
"ml" => (
"ocaml",
tree_sitter_ocaml::LANGUAGE_OCAML.into(),
borrowed(tree_sitter_ocaml::TAGS_QUERY),
),
"mli" => (
"ocaml",
tree_sitter_ocaml::LANGUAGE_OCAML_INTERFACE.into(),
borrowed(tree_sitter_ocaml::TAGS_QUERY),
),
"ex" | "exs" => (
"elixir",
tree_sitter_elixir::LANGUAGE.into(),
borrowed(tree_sitter_elixir::TAGS_QUERY),
),
"sh" | "bash" => (
"bash",
tree_sitter_bash::LANGUAGE.into(),
borrowed(include_str!("queries/bash/tags.scm")),
),
"sql" => (
"sql",
tree_sitter_sequel::LANGUAGE.into(),
borrowed(include_str!("queries/sql/tags.scm")),
),
_ => return None,
};
let grammar_key = match ext {
"mli" => "ocaml-interface",
_ => lang,
};
Some(TagLang {
lang,
grammar_key,
language,
query,
})
}
type TagConfig = std::sync::Arc<tree_sitter_tags::TagsConfiguration>;
static TAG_CONFIGS: std::sync::LazyLock<
std::sync::Mutex<std::collections::HashMap<&'static str, Option<TagConfig>>>,
> = std::sync::LazyLock::new(|| std::sync::Mutex::new(std::collections::HashMap::new()));
fn tag_config(def: &TagLang) -> Option<TagConfig> {
let mut cache = TAG_CONFIGS
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
cache
.entry(def.grammar_key)
.or_insert_with(|| {
tree_sitter_tags::TagsConfiguration::new(def.language.clone(), &def.query, "")
.ok()
.map(std::sync::Arc::new)
})
.clone()
}
fn import_query_for(lang: &str) -> Option<&'static str> {
Some(match lang {
"python" => {
"(import_statement name: (dotted_name) @path)\n\
(import_statement name: (aliased_import name: (dotted_name) @path))\n\
(import_from_statement module_name: (dotted_name) @path)\n\
(import_from_statement module_name: (relative_import) @path)"
}
"javascript" | "typescript" | "tsx" => {
"(import_statement source: (string (string_fragment) @path))\n\
(export_statement source: (string (string_fragment) @path))"
}
"go" => "(import_spec path: (interpreted_string_literal) @path)",
"java" => {
"(import_declaration (scoped_identifier) @path)\n\
(import_declaration (identifier) @path)"
}
"c" | "cpp" => {
"(preproc_include path: (string_literal) @path)\n\
(preproc_include path: (system_lib_string) @path)"
}
_ => return None,
})
}
type ImportQuery = std::sync::Arc<tree_sitter::Query>;
static IMPORT_QUERIES: std::sync::LazyLock<
std::sync::Mutex<std::collections::HashMap<&'static str, Option<ImportQuery>>>,
> = std::sync::LazyLock::new(|| std::sync::Mutex::new(std::collections::HashMap::new()));
fn import_query(def: &TagLang) -> Option<ImportQuery> {
let mut cache = IMPORT_QUERIES
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
cache
.entry(def.grammar_key)
.or_insert_with(|| {
let src = import_query_for(def.lang)?;
tree_sitter::Query::new(&def.language, src)
.ok()
.map(std::sync::Arc::new)
})
.clone()
}
fn normalize_import(raw: &str) -> String {
raw.trim()
.trim_matches(|c| c == '"' || c == '\'' || c == '<' || c == '>')
.trim()
.to_owned()
}
fn append_import_facts(
path: &str,
def: &TagLang,
bytes: &[u8],
nodes: &mut Vec<Node>,
edges: &mut Vec<Edge>,
) {
use streaming_iterator::StreamingIterator as _;
let Some(query) = import_query(def) else {
return;
};
let mut parser = tree_sitter::Parser::new();
if parser.set_language(&def.language).is_err() {
return;
}
let Some(tree) = parser.parse(bytes, None) else {
return;
};
let mut cursor = tree_sitter::QueryCursor::new();
let mut seen = std::collections::BTreeSet::new();
let mut matches = cursor.matches(&query, tree.root_node(), bytes);
while let Some(m) = matches.next() {
for cap in m.captures {
let Ok(raw) = cap.node.utf8_text(bytes) else {
continue;
};
let module = normalize_import(raw);
if module.is_empty() {
continue;
}
let key = format!("import:{}:{module}", def.lang);
if seen.insert(key.clone()) {
nodes.push(Node {
key: key.clone(),
kind: NodeKind::Other("import".into()),
name: module,
path: None,
lang: Some(def.lang.to_owned()),
blob_hash: None,
span: None,
provenance: Provenance::Derived,
meta: serde_json::Value::Null,
});
edges.push(Edge::derived(file_key(path), key, EdgeKind::Imports));
}
}
}
}
fn tag_node_kind(syntax_type: &str) -> NodeKind {
match syntax_type {
"function" | "method" | "constructor" => NodeKind::Fn,
"class" | "struct" => NodeKind::Struct,
"interface" | "trait" | "protocol" => NodeKind::Trait,
"enum" => NodeKind::Enum,
"module" | "namespace" | "object" => NodeKind::Module,
other => NodeKind::Other(other.to_owned()),
}
}
struct TagDef {
name: String,
kind: NodeKind,
range: std::ops::Range<usize>,
docs: Option<String>,
}
fn tag_facts(
path: &str,
blob_id: &str,
bytes: &[u8],
ext: &str,
ingest: IngestConfig,
) -> Option<FactSet> {
let def = tag_lang_for(ext)?;
let lang = def.lang;
let config = tag_config(&def)?;
let mut ctx = tree_sitter_tags::TagsContext::new();
let (tags, _had_error) = ctx.generate_tags(&config, bytes, None).ok()?;
let mut defs: Vec<TagDef> = Vec::new();
let mut calls: Vec<(usize, String)> = Vec::new();
for tag in tags {
let Ok(tag) = tag else { continue };
let Some(name) = bytes
.get(tag.name_range.clone())
.and_then(|b| std::str::from_utf8(b).ok())
else {
continue;
};
let syntax = config.syntax_type_name(tag.syntax_type_id);
if tag.is_definition {
defs.push(TagDef {
name: name.to_owned(),
kind: tag_node_kind(syntax),
range: tag.range.clone(),
docs: tag.docs.clone(),
});
} else if syntax == "call" || syntax == "send" {
calls.push((tag.range.start, name.to_owned()));
}
}
let parents: Vec<Option<usize>> = (0..defs.len())
.map(|i| smallest_enclosing(&defs, defs[i].range.clone(), Some(i)))
.collect();
let keys: Vec<String> = (0..defs.len())
.map(|i| {
let qualified = qualified_name(&defs, &parents, i);
format!("sym:{lang}:{path}#{qualified}")
})
.collect();
let mut nodes = vec![file_node(path, blob_id, bytes, Some(lang), ingest)];
let mut edges: Vec<Edge> = Vec::new();
for (i, d) in defs.iter().enumerate() {
let mut meta = serde_json::Map::new();
if let Some(doc) = &d.docs {
let content = cap_content(doc);
if !content.is_empty() {
meta.insert("content".into(), serde_json::Value::from(content));
}
}
if d.kind == NodeKind::Fn {
let mut names: Vec<String> = calls
.iter()
.filter(|(off, _)| d.range.contains(off))
.filter(|(off, _)| smallest_enclosing_off(&defs, *off) == Some(i))
.map(|(_, name)| name.clone())
.collect();
names.sort();
names.dedup();
if !names.is_empty() {
meta.insert("calls".into(), serde_json::Value::from(names));
}
}
let start = u32::try_from(d.range.start).unwrap_or(u32::MAX);
let end = u32::try_from(d.range.end).unwrap_or(u32::MAX);
nodes.push(Node {
key: keys[i].clone(),
kind: d.kind.clone(),
name: d.name.clone(),
path: Some(path.to_owned()),
lang: Some(lang.to_owned()),
blob_hash: Some(blob_id.to_owned()),
span: Some(Span::new(start, end)),
provenance: Provenance::Derived,
meta: serde_json::Value::Object(meta),
});
match parents[i] {
Some(p) => edges.push(Edge::derived(
keys[p].clone(),
keys[i].clone(),
EdgeKind::Contains,
)),
None => edges.push(Edge::derived(
file_key(path),
keys[i].clone(),
EdgeKind::Defines,
)),
}
}
append_import_facts(path, &def, bytes, &mut nodes, &mut edges);
nodes.sort_by(|a, b| a.key.cmp(&b.key));
nodes.dedup_by(|a, b| a.key == b.key);
edges.sort_by(|a, b| (a.kind.as_str(), &a.src, &a.dst).cmp(&(b.kind.as_str(), &b.src, &b.dst)));
edges.dedup();
Some(FactSet { nodes, edges })
}
fn smallest_enclosing(
defs: &[TagDef],
range: std::ops::Range<usize>,
skip: Option<usize>,
) -> Option<usize> {
let mut best: Option<usize> = None;
for (j, c) in defs.iter().enumerate() {
if Some(j) == skip {
continue;
}
let encloses = c.range.start <= range.start
&& c.range.end >= range.end
&& (c.range.end - c.range.start) > (range.end - range.start);
if encloses
&& best.is_none_or(|b| {
defs[b].range.end - defs[b].range.start > c.range.end - c.range.start
})
{
best = Some(j);
}
}
best
}
fn smallest_enclosing_off(defs: &[TagDef], off: usize) -> Option<usize> {
let mut best: Option<usize> = None;
for (j, c) in defs.iter().enumerate() {
if c.range.contains(&off)
&& best.is_none_or(|b| {
defs[b].range.end - defs[b].range.start > c.range.end - c.range.start
})
{
best = Some(j);
}
}
best
}
fn qualified_name(defs: &[TagDef], parents: &[Option<usize>], i: usize) -> String {
let mut chain: Vec<&str> = vec![defs[i].name.as_str()];
let mut cur = parents[i];
let mut guard = defs.len();
while let Some(p) = cur {
if guard == 0 {
break;
}
guard -= 1;
chain.push(defs[p].name.as_str());
cur = parents[p];
}
chain.reverse();
chain.join("::")
}
fn span(node: tree_sitter::Node) -> Span {
let start = u32::try_from(node.start_byte()).unwrap_or(u32::MAX);
let end = u32::try_from(node.end_byte()).unwrap_or(u32::MAX);
Span::new(start, end)
}
fn qualify(scope: &[Scope], name: &str) -> String {
let mut parts: Vec<&str> = scope.iter().map(|s| s.seg.as_str()).collect();
parts.push(name);
parts.join("::")
}
fn qualify_callee(qualifier: Option<&str>, name: &str) -> String {
match qualifier {
Some(q) if !q.is_empty() && !matches!(q, "self" | "crate" | "super") => {
format!("{q}::{name}")
}
_ => name.to_owned(),
}
}
fn extend(scope: &[Scope], seg: &str, key: Option<String>) -> Vec<Scope> {
let mut next: Vec<Scope> = scope
.iter()
.map(|s| Scope {
seg: s.seg.clone(),
key: s.key.clone(),
})
.collect();
next.push(Scope {
seg: seg.to_owned(),
key,
});
next
}
#[cfg(test)]
mod tests {
use super::{Extractor, FileNodeExtractor, Registry, RustExtractor};
use crate::{EdgeKind, Node, NodeKind};
#[test]
fn file_node_extractor_is_deterministic_and_tagged() {
let ex = FileNodeExtractor;
let a = ex.extract("src/lib.rs", "abc123", b"one\ntwo\n");
let b = ex.extract("src/lib.rs", "abc123", b"one\ntwo\n");
assert_eq!(a, b, "extraction must be deterministic");
assert_eq!(a.nodes.len(), 1);
assert!(a.edges.is_empty());
let node = &a.nodes[0];
assert_eq!(node.key, "file:src/lib.rs");
assert_eq!(node.kind, NodeKind::File);
assert_eq!(node.name, "lib.rs");
assert_eq!(node.blob_hash.as_deref(), Some("abc123"));
assert_eq!(node.meta["lines"], 2);
assert_eq!(node.meta["bytes"], 8);
}
#[test]
fn config_files_emit_config_key_nodes() {
let reg = Registry::new(crate::IngestConfig::default());
let toml = b"[serve]\naddr = \"0.0.0.0:8443\"\ntools = false\n";
let a = reg.extract("config.toml", "cfg1", toml);
let b = reg.extract("config.toml", "cfg1", toml);
assert_eq!(a, b, "config extraction must be deterministic");
assert!(a.nodes.iter().any(|n| n.key == "file:config.toml"));
let addr = a
.nodes
.iter()
.find(|n| n.key == "cfgkey:config.toml#serve.addr")
.expect("serve.addr config_key node");
assert_eq!(addr.kind, NodeKind::Other("config_key".into()));
assert_eq!(addr.name, "serve.addr");
assert_eq!(addr.meta["value"], "0.0.0.0:8443"); assert!(a.edges.iter().any(|e| {
e.src == "file:config.toml"
&& e.dst == "cfgkey:config.toml#serve.addr"
&& e.kind == EdgeKind::Contains
}));
let env = reg.extract(".env", "env1", b"PORT=8080\nPORT=9090\nAPI_TOKEN=s3cr3t\n");
let port = env
.nodes
.iter()
.find(|n| n.key == "cfgkey:.env#PORT")
.expect("PORT node");
assert_eq!(port.meta["value"], "9090", "dotenv last-one-wins");
assert_eq!(
env.nodes
.iter()
.filter(|n| n.key == "cfgkey:.env#PORT")
.count(),
1
);
let token = env
.nodes
.iter()
.find(|n| n.key == "cfgkey:.env#API_TOKEN")
.expect("API_TOKEN node");
assert_eq!(token.meta["value"], "<redacted>", "secret not persisted");
let rs = reg.extract("src/lib.rs", "x", b"pub fn f() {}\n");
assert!(
rs.nodes
.iter()
.all(|n| n.kind != NodeKind::Other("config_key".into()))
);
}
#[test]
fn dockerfile_emits_image_ref_nodes_and_skips_internal_stages() {
let reg = Registry::new(crate::IngestConfig::default());
let df = b"FROM --platform=linux/amd64 rust:1.90 AS builder\nRUN cargo build\n\
FROM builder AS test\nFROM registry.io/app:1.2@sha256:abc AS run\nFROM scratch\n";
let a = reg.extract("Dockerfile", "d1", df);
let b = reg.extract("Dockerfile", "d1", df);
assert_eq!(a, b, "dockerfile extraction must be deterministic");
let refs: Vec<&Node> = a
.nodes
.iter()
.filter(|n| n.kind == NodeKind::Other("image_ref".into()))
.collect();
assert_eq!(refs.len(), 2, "got: {refs:?}");
let rust = refs
.iter()
.find(|n| n.meta["image"] == "rust")
.expect("rust");
assert_eq!(rust.meta["tag"], "1.90");
let app = refs
.iter()
.find(|n| n.meta["image"] == "registry.io/app:1.2")
.expect("app digest");
assert_eq!(app.meta["digest"], "sha256:abc");
assert!(
a.edges
.iter()
.any(|e| { e.src == "file:Dockerfile" && e.kind == EdgeKind::References })
);
assert!(
reg.extract("Dockerfile.prod", "d2", b"FROM alpine:3\n")
.nodes
.iter()
.any(|n| n.kind == NodeKind::Other("image_ref".into()))
);
let c = reg.extract("Dockerfile", "d3", b"FROM alpine AS alpine\n");
assert!(
c.nodes
.iter()
.any(|n| n.kind == NodeKind::Other("image_ref".into())
&& n.meta["image"] == "alpine"),
"FROM x AS x is an external pin, got: {:?}",
c.nodes
);
}
const SAMPLE: &str = r"
use std::path::Path;
pub struct Store;
impl Store {
pub fn open() -> Store {
helper();
Store
}
}
fn helper() {}
mod inner {
pub fn nested() {}
}
";
fn keys(fs: &crate::FactSet) -> Vec<String> {
let mut k: Vec<_> = fs.nodes.iter().map(|n| n.key.clone()).collect();
k.sort();
k
}
#[test]
fn rust_extractor_emits_symbols_and_edges() {
let fs = RustExtractor.extract("src/lib.rs", "blob1", SAMPLE.as_bytes());
let ks = keys(&fs);
assert!(ks.contains(&"file:src/lib.rs".to_owned()));
assert!(ks.contains(&"sym:rust:src/lib.rs#Store".to_owned()));
assert!(ks.contains(&"sym:rust:src/lib.rs#Store::open".to_owned()));
assert!(ks.contains(&"sym:rust:src/lib.rs#helper".to_owned()));
assert!(ks.contains(&"sym:rust:src/lib.rs#inner".to_owned()));
assert!(ks.contains(&"sym:rust:src/lib.rs#inner::nested".to_owned()));
let open = fs
.nodes
.iter()
.find(|n| n.key == "sym:rust:src/lib.rs#Store::open")
.expect("open node");
assert_eq!(open.meta["calls"], serde_json::json!(["helper"]));
let defines: Vec<_> = fs
.edges
.iter()
.filter(|e| e.kind == EdgeKind::Defines && e.dst == "sym:rust:src/lib.rs#helper")
.collect();
assert_eq!(defines.len(), 1);
assert!(fs.edges.iter().any(|e| e.kind == EdgeKind::Contains
&& e.src == "sym:rust:src/lib.rs#inner"
&& e.dst == "sym:rust:src/lib.rs#inner::nested"));
assert!(fs.edges.iter().any(|e| e.kind == EdgeKind::Imports
&& e.src == "file:src/lib.rs"
&& e.dst == "import:rust:std::path::Path"));
}
#[test]
fn rust_extractor_records_struct_field_names() {
let src = "pub struct ServeConfig {\n\
\x20 pub addr: Option<String>,\n\
\x20 pub tls_cert: Option<String>,\n\
}\n\
pub struct Pair(u8, u8);\n\
pub struct Marker;\n";
let fs = RustExtractor.extract("src/config.rs", "b", src.as_bytes());
let fields = |key: &str| {
fs.nodes
.iter()
.find(|n| n.key == key)
.and_then(|n| n.meta.get("fields").cloned())
};
assert_eq!(
fields("sym:rust:src/config.rs#ServeConfig"),
Some(serde_json::json!(["addr", "tls_cert"])),
"named fields captured in source order"
);
assert_eq!(fields("sym:rust:src/config.rs#Pair"), None);
assert_eq!(fields("sym:rust:src/config.rs#Marker"), None);
}
#[test]
fn struct_records_field_types_and_config_root_marker() {
let src = "// @rto:config\n\
pub struct Config {\n\
\x20 pub zerobus: ZerobusConfig,\n\
\x20 pub replicas: Option<u32>,\n\
}\n\
pub struct ZerobusConfig {\n\
\x20 pub server_endpoint: String,\n\
}\n";
let fs = RustExtractor.extract("src/config.rs", "b", src.as_bytes());
let node = |key: &str| fs.nodes.iter().find(|n| n.key == key).expect("node");
let root = node("sym:rust:src/config.rs#Config");
assert_eq!(root.meta.get("config_root"), Some(&serde_json::json!(true)));
assert_eq!(
root.meta.get("field_types"),
Some(&serde_json::json!({ "zerobus": "ZerobusConfig", "replicas": "u32" })),
"transparent wrappers peeled (Option<u32> → u32)"
);
assert_eq!(
node("sym:rust:src/config.rs#ZerobusConfig")
.meta
.get("config_root"),
None
);
}
#[test]
fn config_root_struct_synthesizes_recursive_dotted_config_keys() {
let src = "// @rto:config\n\
pub struct Config {\n\
\x20 pub zerobus: ZerobusConfig,\n\
\x20 pub log_level: String,\n\
}\n\
pub struct ZerobusConfig {\n\
\x20 pub server_endpoint: String,\n\
\x20 pub workspace_url: String,\n\
}\n";
let fs = RustExtractor.extract("src/config.rs", "b", src.as_bytes());
let cfg = |dotted: &str| {
fs.nodes
.iter()
.find(|n| n.key == format!("cfgkey:src/config.rs#{dotted}"))
};
for dotted in [
"zerobus.server_endpoint",
"zerobus.workspace_url",
"log_level",
] {
let n = cfg(dotted).unwrap_or_else(|| panic!("missing {dotted}: {:?}", fs.nodes));
assert_eq!(n.kind, NodeKind::Other("config_key".into()));
assert_eq!(n.meta.get("key").and_then(|v| v.as_str()), Some(dotted));
assert_eq!(
n.meta.get("source").and_then(|v| v.as_str()),
Some("struct")
);
assert_eq!(
n.meta.get("struct").and_then(|v| v.as_str()),
Some("Config")
);
}
assert!(
cfg("zerobus").is_none(),
"intermediate section is not a leaf"
);
assert!(fs.edges.iter().any(|e| e.src == "file:src/config.rs"
&& e.dst == "cfgkey:src/config.rs#zerobus.server_endpoint"
&& e.kind == EdgeKind::Contains));
}
#[test]
fn struct_without_config_marker_synthesizes_no_config_keys() {
let src = "pub struct Config {\n\
\x20 pub zerobus: ZerobusConfig,\n\
}\n\
pub struct ZerobusConfig {\n\
\x20 pub server_endpoint: String,\n\
}\n";
let fs = RustExtractor.extract("src/config.rs", "b", src.as_bytes());
assert!(
fs.nodes
.iter()
.all(|n| n.kind != NodeKind::Other("config_key".into())),
"no synthetic config_key nodes without the marker: {:?}",
fs.nodes
);
}
#[test]
fn config_root_recursion_terminates_on_a_type_cycle() {
let src = "// @rto:config\n\
pub struct Config {\n\
\x20 pub addr: String,\n\
\x20 pub next: Box<Config>,\n\
}\n";
let fs = RustExtractor.extract("src/config.rs", "b", src.as_bytes());
let has = |dotted: &str| {
fs.nodes
.iter()
.any(|n| n.key == format!("cfgkey:src/config.rs#{dotted}"))
};
assert!(has("addr"));
assert!(has("next"), "cyclic field falls back to a leaf");
assert!(!has("next.addr"), "no unbounded expansion");
}
#[test]
fn rust_extraction_is_deterministic() {
let a = RustExtractor.extract("src/lib.rs", "blob1", SAMPLE.as_bytes());
let b = RustExtractor.extract("src/lib.rs", "blob1", SAMPLE.as_bytes());
assert_eq!(a, b);
}
#[test]
fn rust_extractor_captures_doc_comments() {
let src = "/// The central store.\n\
pub struct Store;\n\n\
/// Opens it.\n\
/// Reads the config.\n\
pub fn open() {}\n\n\
// not a doc comment\n\
pub fn plain() {}\n";
let fs = RustExtractor.extract("src/lib.rs", "b", src.as_bytes());
let content = |key: &str| {
fs.nodes
.iter()
.find(|n| n.key == key)
.and_then(|n| n.meta.get("content"))
.and_then(|v| v.as_str())
.map(ToOwned::to_owned)
};
assert_eq!(
content("sym:rust:src/lib.rs#Store").as_deref(),
Some("The central store.")
);
assert_eq!(
content("sym:rust:src/lib.rs#open").as_deref(),
Some("Opens it. Reads the config.")
);
assert_eq!(content("sym:rust:src/lib.rs#plain"), None);
}
#[test]
fn prose_file_captures_capped_body() {
let md = FileNodeExtractor.extract("docs/x.md", "b", b"# Title\n\nSome prose here.\n");
assert_eq!(md.nodes[0].meta["content"], "# Title Some prose here.");
let rs = FileNodeExtractor.extract("notes.bin", "b", b"\x00\x01binary");
assert!(rs.nodes[0].meta.get("content").is_none());
let upper = FileNodeExtractor.extract("README.MD", "b", b"# Hi\n");
assert_eq!(upper.nodes[0].meta["content"], "# Hi");
}
#[cfg(feature = "pdf-text")]
fn minimal_pdf(text: &str) -> Vec<u8> {
let content = format!("BT /F1 24 Tf 72 720 Td ({text}) Tj ET");
let objects = [
"<< /Type /Catalog /Pages 2 0 R >>".to_owned(),
"<< /Type /Pages /Kids [3 0 R] /Count 1 >>".to_owned(),
"<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] /Contents 4 0 R /Resources << /Font << /F1 5 0 R >> >> >>".to_owned(),
format!("<< /Length {} >>\nstream\n{content}\nendstream", content.len()),
"<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>".to_owned(),
];
let mut pdf = Vec::new();
pdf.extend_from_slice(b"%PDF-1.4\n");
let mut offsets = Vec::new();
for (i, obj) in objects.iter().enumerate() {
offsets.push(pdf.len());
pdf.extend_from_slice(format!("{} 0 obj\n{obj}\nendobj\n", i + 1).as_bytes());
}
let xref_start = pdf.len();
pdf.extend_from_slice(
format!("xref\n0 {}\n0000000000 65535 f \n", objects.len() + 1).as_bytes(),
);
for off in &offsets {
pdf.extend_from_slice(format!("{off:010} 00000 n \n").as_bytes());
}
pdf.extend_from_slice(
format!(
"trailer\n<< /Size {} /Root 1 0 R >>\nstartxref\n{xref_start}\n%%EOF\n",
objects.len() + 1
)
.as_bytes(),
);
pdf
}
#[cfg(feature = "pdf-text")]
#[test]
fn pdf_file_captures_text_content() {
let pdf = minimal_pdf("Hello Roteiro");
let facts = FileNodeExtractor.extract("docs/guide.pdf", "b", &pdf);
let content = facts.nodes[0].meta["content"].as_str().unwrap();
assert!(content.contains("Hello Roteiro"), "got: {content:?}");
let upper = FileNodeExtractor.extract("docs/Guide.PDF", "b", &pdf);
assert!(upper.nodes[0].meta.get("content").is_some());
let bad = FileNodeExtractor.extract("docs/bad.pdf", "b", b"%PDF-1.4\ngarbage");
assert!(bad.nodes[0].meta.get("content").is_none());
}
#[cfg(any(feature = "image-ocr", feature = "image-vision"))]
#[test]
fn image_content_guards_before_touching_models() {
use crate::media::{MAX_IMAGE_BYTES, is_image};
assert!(is_image("shot.PNG"));
assert!(is_image("b.jpeg"));
assert!(is_image("c.jpg"));
assert!(!is_image("d.gif"));
assert!(
super::image_content("notes.txt", b"hello", super::IngestConfig::default()).is_none()
);
let big = vec![0u8; MAX_IMAGE_BYTES + 1];
assert!(super::image_content("shot.png", &big, super::IngestConfig::default()).is_none());
}
#[test]
fn doc_comment_body_recognises_doc_markers() {
assert_eq!(super::doc_comment_body("/// hi").as_deref(), Some("hi"));
assert_eq!(
super::doc_comment_body("//! mod doc").as_deref(),
Some("mod doc")
);
assert_eq!(
super::doc_comment_body("/** block */").as_deref(),
Some("block")
);
assert_eq!(super::doc_comment_body("// plain"), None);
assert_eq!(super::doc_comment_body("//// header"), None);
assert_eq!(super::doc_comment_body("/**/").as_deref(), Some(""));
assert_eq!(super::doc_comment_body("/*!*/").as_deref(), Some(""));
}
#[test]
fn registry_dispatches_by_extension() {
let rs = Registry::default().extract("src/lib.rs", "b", SAMPLE.as_bytes());
assert!(rs.nodes.len() > 1, "rust file yields symbols");
let txt = Registry::default().extract("notes.txt", "b", b"hello\n");
assert_eq!(
txt.nodes.len(),
1,
"non-code file falls back to a file node"
);
assert_eq!(txt.nodes[0].kind, NodeKind::File);
}
#[test]
fn tags_extracts_python_symbols_calls_and_nesting() {
let src = "def helper():\n pass\n\nclass Thing:\n def run(self):\n helper()\n";
let fs = Registry::default().extract("app.py", "b", src.as_bytes());
let names: Vec<&str> = fs.nodes.iter().map(|n| n.name.as_str()).collect();
assert!(names.contains(&"helper"), "top-level function");
assert!(names.contains(&"Thing"), "class");
assert!(names.contains(&"run"), "method");
assert_eq!(
fs.nodes
.iter()
.find(|n| n.name == "helper")
.and_then(|n| n.lang.as_deref()),
Some("python")
);
assert!(
fs.edges
.iter()
.any(|e| e.kind == EdgeKind::Contains && e.dst.ends_with("#Thing::run")),
"method nested under class via containment"
);
let run = fs.nodes.iter().find(|n| n.name == "run").unwrap();
let calls = run.meta.get("calls").and_then(|v| v.as_array()).unwrap();
assert!(
calls.iter().any(|c| c.as_str() == Some("helper")),
"enclosed call captured in meta.calls"
);
}
#[test]
fn tags_extraction_is_deterministic() {
let src = b"package main\nfunc Add(a int) int { return a }\n";
let a = Registry::default().extract("m.go", "b", src);
let b = Registry::default().extract("m.go", "b", src);
assert_eq!(a, b, "tags extraction must be deterministic");
assert!(
a.nodes
.iter()
.any(|n| n.name == "Add" && n.kind == NodeKind::Fn)
);
}
#[test]
fn tags_extracts_typescript() {
let ts = Registry::default().extract("svc.ts", "b", b"export class Svc {\n run() {}\n}\n");
assert!(ts.nodes.iter().any(|n| n.name == "Svc"), "class");
assert!(ts.nodes.iter().any(|n| n.name == "run"), "method");
assert_eq!(
ts.nodes
.iter()
.find(|n| n.name == "Svc")
.and_then(|n| n.lang.as_deref()),
Some("typescript")
);
}
fn import_targets(path: &str, src: &[u8]) -> Vec<String> {
Registry::default()
.extract(path, "b", src)
.nodes
.iter()
.filter(|n| n.kind == NodeKind::Other("import".into()))
.inspect(|n| {
assert!(
n.path.is_none(),
"import node must not be file-scoped: {}",
n.key
);
})
.map(|n| n.key.clone())
.collect()
}
#[test]
fn extracts_imports_edges_per_language() {
let cases: &[(&str, &[u8], &[&str])] = &[
(
"app.py",
b"import os\nfrom a.b import c\nimport x.y as z\n",
&["import:python:os", "import:python:a.b", "import:python:x.y"],
),
(
"m.js",
b"import foo from \"./mod.js\";\nexport { y } from \"./y.js\";\n",
&["import:javascript:./mod.js", "import:javascript:./y.js"],
),
(
"svc.ts",
b"import { A } from \"./a\";\n",
&["import:typescript:./a"],
),
(
"m.go",
b"package main\nimport (\n\t\"fmt\"\n\t\"os\"\n)\n",
&["import:go:fmt", "import:go:os"],
),
(
"M.java",
b"import java.util.List;\nimport static a.B.c;\n",
&["import:java:java.util.List", "import:java:a.B.c"],
),
(
"m.c",
b"#include <stdio.h>\n#include \"local.h\"\n",
&["import:c:stdio.h", "import:c:local.h"],
),
("m.cpp", b"#include <vector>\n", &["import:cpp:vector"]),
];
for (path, src, expected) in cases {
let got = import_targets(path, src);
for want in *expected {
assert!(
got.iter().any(|k| k == want),
"{path}: expected import node {want}, got {got:?}"
);
}
let fs = Registry::default().extract(path, "b", src);
for want in *expected {
assert!(
fs.edges.iter().any(|e| e.kind == EdgeKind::Imports
&& e.src == format!("file:{path}")
&& &e.dst == want),
"{path}: expected Imports edge to {want}"
);
}
}
}
#[test]
fn every_registered_language_query_compiles() {
for ext in [
"py", "js", "ts", "tsx", "go", "rb", "java", "c", "cpp", "cs", "php", "scala", "ml",
"mli", "ex", "sh", "sql",
] {
let def = super::tag_lang_for(ext).unwrap_or_else(|| panic!("no language for .{ext}"));
let lang = def.lang;
assert!(
super::tag_config(&def).is_some(),
"tags query for .{ext} ({lang}) must compile against its grammar"
);
}
}
#[test]
fn ocaml_impl_and_interface_cache_under_distinct_grammars() {
let ml = super::tag_lang_for("ml").unwrap();
let mli = super::tag_lang_for("mli").unwrap();
assert_eq!(ml.lang, "ocaml");
assert_eq!(mli.lang, "ocaml");
assert_ne!(
ml.grammar_key, mli.grammar_key,
"distinct grammars must cache separately"
);
}
#[test]
fn tags_extracts_vendored_bash_query() {
let src = "greet() {\n echo hi\n}\nmain() {\n greet\n}\n";
let fs = Registry::default().extract("run.sh", "b", src.as_bytes());
let names: Vec<&str> = fs.nodes.iter().map(|n| n.name.as_str()).collect();
assert!(names.contains(&"greet"), "shell function greet");
assert!(names.contains(&"main"), "shell function main");
let main = fs.nodes.iter().find(|n| n.name == "main").unwrap();
assert!(
main.meta
.get("calls")
.and_then(|v| v.as_array())
.is_some_and(|c| c.iter().any(|x| x.as_str() == Some("greet"))),
"internal command invocation captured"
);
}
#[test]
fn tags_extracts_vendored_sql_query() {
let src = "CREATE TABLE users (id int);\n\
CREATE FUNCTION recent() RETURNS int AS $$ SELECT total(id) FROM users $$ LANGUAGE sql;\n";
let fs = Registry::default().extract("schema.sql", "b", src.as_bytes());
let names: Vec<&str> = fs.nodes.iter().map(|n| n.name.as_str()).collect();
assert!(names.contains(&"users"), "table definition");
assert!(names.contains(&"recent"), "function definition");
assert_eq!(
fs.nodes.iter().find(|n| n.name == "users").map(|n| &n.kind),
Some(&NodeKind::Other("table".to_owned()))
);
let f = fs.nodes.iter().find(|n| n.name == "recent").unwrap();
assert!(
f.meta
.get("calls")
.and_then(|v| v.as_array())
.is_some_and(|c| c.iter().any(|x| x.as_str() == Some("total"))),
"invocation inside function captured in meta.calls"
);
assert_eq!(
fs.nodes
.iter()
.find(|n| n.name == "users")
.and_then(|n| n.lang.as_deref()),
Some("sql")
);
}
#[test]
fn ingest_prose_toggle_gates_embedded_content() {
use super::IngestConfig;
let content = |ingest: IngestConfig| {
Registry::new(ingest)
.extract("notes.md", "b", b"# Title\n\nBody text.\n")
.nodes[0]
.meta
.get("content")
.and_then(|v| v.as_str())
.map(str::to_owned)
};
assert!(
content(IngestConfig::default()).is_some_and(|c| c.contains("Body text")),
"prose content embedded by default"
);
assert_eq!(
content(IngestConfig {
prose: false,
..IngestConfig::default()
}),
None,
"disabling prose suppresses the embedded body"
);
}
#[test]
fn env_tag_stable_by_default_and_shifts_when_gated() {
use super::IngestConfig;
let all_on = Registry::new(IngestConfig::default()).env_tag();
assert_eq!(all_on, Registry::default().env_tag());
let no_prose = Registry::new(IngestConfig {
prose: false,
..IngestConfig::default()
})
.env_tag();
let no_pdf = Registry::new(IngestConfig {
pdf: false,
..IngestConfig::default()
})
.env_tag();
let no_ocr = Registry::new(IngestConfig {
ocr: false,
..IngestConfig::default()
})
.env_tag();
assert_ne!(no_prose, all_on);
assert_ne!(no_pdf, all_on);
assert_ne!(no_ocr, all_on);
assert_ne!(no_prose, no_pdf);
assert_ne!(no_ocr, no_prose);
assert_ne!(no_ocr, no_pdf);
}
#[test]
fn generation_toggles_do_not_move_the_extraction_cache_key() {
use super::IngestConfig;
let all_on = Registry::default().env_tag();
for (label, cfg) in [
(
"audio",
IngestConfig {
audio: false,
..IngestConfig::default()
},
),
(
"vision",
IngestConfig {
vision: false,
..IngestConfig::default()
},
),
(
"both",
IngestConfig {
audio: false,
vision: false,
..IngestConfig::default()
},
),
] {
assert_eq!(
Registry::new(cfg).env_tag(),
all_on,
"`{label}` gates generation, not extraction, so it must not move the cache key",
);
}
}
#[test]
fn generation_toggles_gate_media_build() {
use super::IngestConfig;
use crate::media::MediaKind;
let all_on = IngestConfig::default();
assert!(all_on.generates(MediaKind::Audio));
assert!(all_on.generates(MediaKind::Vision));
let no_audio = IngestConfig {
audio: false,
..IngestConfig::default()
};
assert!(!no_audio.generates(MediaKind::Audio));
assert!(
no_audio.generates(MediaKind::Vision),
"each modality is gated independently"
);
}
}
#[cfg(all(test, feature = "image-vision"))]
fn tiny_png() -> Vec<u8> {
let img = image::RgbImage::from_fn(32, 32, |x, y| {
if x == y {
image::Rgb([0, 0, 0])
} else {
image::Rgb([255, 255, 255])
}
});
let mut png = std::io::Cursor::new(Vec::new());
image::DynamicImage::ImageRgb8(img)
.write_to(&mut png, image::ImageFormat::Png)
.expect("encode png");
png.into_inner()
}
#[cfg(all(test, any(feature = "image-vision", feature = "audio-transcribe")))]
fn serialise_media_engine_test() -> std::sync::MutexGuard<'static, ()> {
static LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
LOCK.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
#[cfg(all(test, feature = "image-vision"))]
mod vision_engine_teardown {
use super::{release_media_engines, serialise_media_engine_test, tiny_png};
use crate::media::producers::{VLM_MODEL, vlm_content};
#[test]
fn describing_an_image_leaves_a_releasable_engine() {
let _serial = serialise_media_engine_test();
let dir = crate::models::model_dir(VLM_MODEL);
if !dir.join("model.gguf").exists() || !dir.join("mmproj.gguf").exists() {
eprintln!("SKIP: `{VLM_MODEL}` not installed (run `roteiro model pull {VLM_MODEL}`)");
return;
}
let _description = vlm_content(&tiny_png());
assert!(
release_media_engines(),
"the engine `vlm_content` cached must be released, not leaked to exit"
);
assert!(
!release_media_engines(),
"releasing again must be a no-op, so every exit path can call it"
);
}
}
#[cfg(all(test, feature = "image-vision", feature = "audio-transcribe"))]
mod two_modality_teardown {
use super::{release_media_engines, serialise_media_engine_test, tiny_png};
use crate::media::producers::{
ASR_MODEL, VLM_MODEL, asr_content, asr_engine, vlm_content, vlm_engine,
};
pub(super) const TINY_WAV: &[u8] =
include_bytes!("../tests/fixtures/audio/syllables-16khz-mono-512ms.wav");
fn installed(name: &str) -> bool {
let dir = crate::models::model_dir(name);
dir.join("model.gguf").exists() && dir.join("mmproj.gguf").exists()
}
#[test]
fn both_modalities_get_a_working_engine_in_one_process() {
let _serial = serialise_media_engine_test();
if !installed(VLM_MODEL) || !installed(ASR_MODEL) {
eprintln!(
"SKIP: need both `{VLM_MODEL}` and `{ASR_MODEL}` installed \
(run `roteiro model pull <name>`)"
);
return;
}
assert!(vlm_engine().is_some(), "the vision engine must build");
assert!(
asr_engine().is_some(),
"the second engine must share the first's backend, not be inert (#296)"
);
let png = tiny_png();
let _description = vlm_content(&png);
let _transcript = asr_content(TINY_WAV);
let _description_again = vlm_content(&png);
let _transcript_again = asr_content(TINY_WAV);
let (vision, audio) = (
vlm_engine().expect("resident").projector_inits(),
asr_engine().expect("resident").projector_inits(),
);
assert_eq!(vision, 1, "two images must load the vision projector once");
assert_eq!(audio, 1, "two clips must load the audio projector once");
assert!(
release_media_engines(),
"two engines and a backend must all be released, not leaked to exit"
);
assert!(
!release_media_engines(),
"releasing again must be a no-op, so every exit path can call it"
);
}
}
#[cfg(all(test, feature = "image-vision", feature = "audio-transcribe"))]
mod projector_binding {
use super::two_modality_teardown::TINY_WAV;
use super::{release_media_engines, serialise_media_engine_test, tiny_png};
use crate::media::producers::{ASR_MODEL, VLM_MODEL};
use rto_llama::llama::{LlamaEngine, Served};
use rto_llama::{ChatRequest, Engine, Message};
fn served(name: &str) -> Option<Served> {
let dir = crate::models::model_dir(name);
let (gguf, mmproj) = (dir.join("model.gguf"), dir.join("mmproj.gguf"));
(gguf.exists() && mmproj.exists()).then(|| Served {
name: name.to_owned(),
path: gguf,
mmproj: Some(mmproj),
})
}
fn media_chat(
engine: &LlamaEngine,
model: &str,
images: Vec<Vec<u8>>,
audio: Vec<Vec<u8>>,
) -> String {
engine
.chat(&ChatRequest {
model: model.to_owned(),
messages: vec![Message {
role: "user".to_owned(),
content: "Describe what you perceive in one short sentence.".to_owned(),
}],
images,
audio,
temperature: 0.0,
max_tokens: 32,
})
.expect("the blob reaches its projector and completes")
.content
}
#[test]
fn evicting_a_model_rebuilds_its_projector_rather_than_reusing_a_stale_one() {
let _serial = serialise_media_engine_test();
let (Some(vlm), Some(asr)) = (served(VLM_MODEL), served(ASR_MODEL)) else {
eprintln!(
"SKIP: need both `{VLM_MODEL}` and `{ASR_MODEL}` installed \
(run `roteiro model pull <name>`)"
);
return;
};
let engine = LlamaEngine::new(vec![vlm, asr], 0).expect("engine builds");
let first = media_chat(&engine, ASR_MODEL, Vec::new(), vec![TINY_WAV.to_vec()]);
assert_eq!(engine.projector_inits(), 1, "the audio projector loaded");
let described = media_chat(&engine, VLM_MODEL, vec![tiny_png()], Vec::new());
assert!(
!described.trim().is_empty(),
"a second, different projector must work in the same process (#298)"
);
assert_eq!(
engine.projector_inits(),
2,
"a different mmproj is a different projector — never the first one reused"
);
let again = media_chat(&engine, ASR_MODEL, Vec::new(), vec![TINY_WAV.to_vec()]);
assert_eq!(
engine.projector_inits(),
3,
"a reloaded model gets a freshly bound projector, not the evicted model's"
);
assert_eq!(
first, again,
"and the rebuilt projector produces exactly what the original did"
);
drop(engine);
assert!(
release_media_engines(),
"the backend is releasable once the engine holding it is gone"
);
}
}