use super::display::{print_tool_call, print_tool_output};
use async_trait::async_trait;
use std::sync::Mutex;
use std::time::Duration;
#[async_trait]
pub trait Embedder: Send + Sync {
async fn embed(&self, text: &str) -> anyhow::Result<Vec<f32>>;
}
pub struct EmbeddingsClient {
url: String,
model: String,
kind: crate::BackendKind,
api_key: Option<String>,
timeout_secs: u64,
retries: u32,
}
impl EmbeddingsClient {
pub fn new(
url: impl Into<String>,
model: impl Into<String>,
kind: crate::BackendKind,
api_key: Option<String>,
timeout_secs: u64,
retries: u32,
) -> Self {
Self {
url: url.into(),
model: model.into(),
kind,
api_key,
timeout_secs,
retries,
}
}
pub async fn embed_batch(&self, texts: &[String]) -> anyhow::Result<Vec<Vec<f32>>> {
let mut out = Vec::with_capacity(texts.len());
for t in texts {
out.push(self.embed(t).await?);
}
Ok(out)
}
async fn embed_once(&self, text: &str) -> anyhow::Result<Vec<f32>> {
let base = self.url.trim_end_matches('/');
let (endpoint, body) = match self.kind {
crate::BackendKind::Ollama => (
format!("{base}/api/embeddings"),
serde_json::json!({ "model": self.model, "prompt": text }),
),
crate::BackendKind::Openai => (
format!("{base}/v1/embeddings"),
serde_json::json!({ "model": self.model, "input": text }),
),
crate::BackendKind::Embedded => anyhow::bail!(
"the embedded backend is chat-only and does not serve embeddings; \
set `embeddings_api` to an ollama/openai backend"
),
};
let client = reqwest::Client::builder()
.timeout(Duration::from_secs(self.timeout_secs))
.build()?;
let mut req = client.post(&endpoint).json(&body);
if let Some(key) = &self.api_key {
req = req.bearer_auth(key);
}
let resp = req.send().await?;
if !resp.status().is_success() {
anyhow::bail!("embeddings endpoint {endpoint} returned {}", resp.status());
}
let json: serde_json::Value = resp.json().await?;
let arr = match self.kind {
crate::BackendKind::Ollama => json["embedding"].as_array(),
crate::BackendKind::Openai => json["data"][0]["embedding"].as_array(),
crate::BackendKind::Embedded => None,
}
.ok_or_else(|| anyhow::anyhow!("embeddings response missing `embedding` array"))?;
let vec: Vec<f32> = arr
.iter()
.map(|v| v.as_f64().unwrap_or(0.0) as f32)
.collect();
if vec.is_empty() {
anyhow::bail!("embeddings response had an empty vector");
}
Ok(vec)
}
}
#[async_trait]
impl Embedder for EmbeddingsClient {
async fn embed(&self, text: &str) -> anyhow::Result<Vec<f32>> {
let mut last_err = None;
for attempt in 0..=self.retries {
if attempt > 0 {
let backoff = Duration::from_millis(250u64 << (attempt - 1).min(4));
tokio::time::sleep(backoff).await;
}
match self.embed_once(text).await {
Ok(v) => return Ok(v),
Err(e) => last_err = Some(e),
}
}
Err(last_err.unwrap_or_else(|| anyhow::anyhow!("embeddings failed")))
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CodeChunk {
pub file: String,
pub start_line: usize,
pub end_line: usize,
pub kind: String,
pub text: String,
}
pub const CHUNK_MAX_CHARS: usize = 2_000;
const WINDOW_LINES: usize = 40;
pub fn chunk_source(file: &str, source: &str) -> Vec<CodeChunk> {
let lines: Vec<&str> = source.lines().collect();
if lines.is_empty() {
return Vec::new();
}
let mut defs = crate::symbols::Lang::from_path(file)
.map(|lang| crate::symbols::extract_definitions(source, lang))
.unwrap_or_default();
if defs.is_empty() {
return window_chunks(file, &lines, 1, lines.len());
}
defs.sort_by_key(|d| d.line);
let starts: Vec<usize> = defs.iter().map(|d| block_start(&lines, d.line)).collect();
let mut chunks = Vec::new();
for (i, d) in defs.iter().enumerate() {
let start = starts[i];
let end = if i + 1 < defs.len() {
starts[i + 1].saturating_sub(1).max(start)
} else {
lines.len()
};
let kind = format!("{:?}", d.kind).to_lowercase();
let text = join_lines(&lines, start, end);
if text.chars().count() > CHUNK_MAX_CHARS {
chunks.extend(window_chunks(file, &lines, start, end));
} else {
chunks.push(CodeChunk {
file: file.to_string(),
start_line: start,
end_line: end,
kind,
text,
});
}
}
chunks
}
fn block_start(lines: &[&str], def_line: usize) -> usize {
let is_doc = |s: &str| {
let t = s.trim_start();
t.starts_with("///")
|| t.starts_with("//")
|| t.starts_with('#')
|| t.starts_with("\"\"\"")
|| t.starts_with("/*")
|| t.starts_with('*')
};
let mut start = def_line;
while start > 1 && is_doc(lines[start - 2]) {
start -= 1;
}
start
}
fn join_lines(lines: &[&str], start: usize, end: usize) -> String {
let end = end.min(lines.len());
if start > end {
return String::new();
}
lines[start - 1..end].join("\n")
}
fn window_chunks(file: &str, lines: &[&str], from: usize, to: usize) -> Vec<CodeChunk> {
let mut chunks = Vec::new();
let mut s = from;
while s <= to {
let e = (s + WINDOW_LINES - 1).min(to);
chunks.push(CodeChunk {
file: file.to_string(),
start_line: s,
end_line: e,
kind: "window".to_string(),
text: join_lines(lines, s, e),
});
s = e + 1;
}
chunks
}
pub(crate) const CODE_EVIDENCE_CAP: usize = 6_000;
pub trait SemanticIndex: Send + Sync {
fn index_chunk(&self, chunk: CodeChunk, embedding: Vec<f32>);
fn search(&self, query: &[f32], top_k: usize) -> Vec<(f32, CodeChunk)>;
fn chunks_indexed(&self) -> u64;
fn indexed_chars(&self) -> u64;
fn clear(&self);
}
#[derive(Default)]
pub struct SessionSemanticIndex {
entries: Mutex<Vec<(CodeChunk, Vec<f32>)>>,
}
impl SemanticIndex for SessionSemanticIndex {
fn index_chunk(&self, chunk: CodeChunk, embedding: Vec<f32>) {
self.entries.lock().unwrap().push((chunk, embedding));
}
fn search(&self, query: &[f32], top_k: usize) -> Vec<(f32, CodeChunk)> {
let entries = self.entries.lock().unwrap();
let mut scored: Vec<(f32, CodeChunk)> = entries
.iter()
.map(|(c, e)| (cosine(query, e), c.clone()))
.collect();
scored.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal));
scored.truncate(top_k);
scored
}
fn chunks_indexed(&self) -> u64 {
self.entries.lock().unwrap().len() as u64
}
fn indexed_chars(&self) -> u64 {
self.entries
.lock()
.unwrap()
.iter()
.map(|(c, _)| c.text.chars().count() as u64)
.sum()
}
fn clear(&self) {
self.entries.lock().unwrap().clear();
}
}
pub fn cosine(a: &[f32], b: &[f32]) -> f32 {
if a.is_empty() || a.len() != b.len() {
return 0.0;
}
let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum();
let na = a.iter().map(|x| x * x).sum::<f32>().sqrt();
let nb = b.iter().map(|x| x * x).sum::<f32>().sqrt();
if na == 0.0 || nb == 0.0 {
0.0
} else {
dot / (na * nb)
}
}
fn render_hits(hits: &[(f32, CodeChunk)], total_cap: usize) -> Option<String> {
if hits.is_empty() {
return None;
}
let mut body = String::from("<code_evidence>\n");
for (score, chunk) in hits {
let piece = format!(
"// {}:{}-{} ({}, score {score:.2})\n{}\n\n",
chunk.file, chunk.start_line, chunk.end_line, chunk.kind, chunk.text
);
if body.chars().count() + piece.chars().count() + "</code_evidence>".len() > total_cap {
body.push_str("[… more evidence omitted to fit the budget …]\n");
break;
}
body.push_str(&piece);
}
body.push_str("</code_evidence>");
Some(body)
}
pub(crate) fn build_code_evidence_block(
index: &dyn SemanticIndex,
query: &[f32],
top_k: usize,
total_cap: usize,
) -> Option<String> {
render_hits(&index.search(query, top_k), total_cap)
}
pub fn code_evidence_block(
index: &dyn SemanticIndex,
query: &[f32],
top_k: usize,
) -> Option<String> {
build_code_evidence_block(index, query, top_k, CODE_EVIDENCE_CAP)
}
pub async fn index_files(
files: &[(String, String)],
embedder: &dyn Embedder,
index: &dyn SemanticIndex,
on_failure: crate::OnEmbedFailure,
) -> usize {
let mut indexed = 0;
for (file, source) in files {
for chunk in chunk_source(file, source) {
match embedder.embed(&chunk.text).await {
Ok(v) => {
index.index_chunk(chunk, v);
indexed += 1;
}
Err(e) => match on_failure {
crate::OnEmbedFailure::Disable => {
tracing::error!(
error = %e,
file = file.as_str(),
"semantic indexing disabled: embeddings failed. Configure \
[context.semantic] for a working embedder: use \
embeddings_endpoint/embeddings_api for an Ollama or OpenAI \
embeddings service, or embeddings_api = \"embedded\" with \
embedding_model_path for local in-process embeddings. Set \
on_embed_failure = \"warn\" to keep trying per-chunk. Indexed \
{indexed} chunk(s) before stopping."
);
return indexed;
}
crate::OnEmbedFailure::Warn => {
tracing::warn!(error = %e, file = file.as_str(), "embed failed; skipping chunk");
}
},
}
}
}
indexed
}
const RERANK_OVERFETCH: usize = 3;
const DEF_BOOST: f32 = 0.05;
const PATH_BOOST: f32 = 0.05;
fn rerank(query: &str, hits: &mut [(f32, CodeChunk)]) {
let terms: Vec<String> = query
.split(|c: char| !c.is_alphanumeric())
.filter(|t| t.len() >= 3)
.map(str::to_lowercase)
.collect();
let boost = |chunk: &CodeChunk| -> f32 {
let mut b = 0.0;
if chunk.kind != "window" {
b += DEF_BOOST;
}
let file_lc = chunk.file.to_lowercase();
if terms.iter().any(|t| file_lc.contains(t.as_str())) {
b += PATH_BOOST;
}
b
};
hits.sort_by(|a, b| {
let sb = b.0 + boost(&b.1);
let sa = a.0 + boost(&a.1);
sb.partial_cmp(&sa).unwrap_or(std::cmp::Ordering::Equal)
});
}
pub async fn retrieve_evidence(
query: &str,
embedder: &dyn Embedder,
index: &dyn SemanticIndex,
top_k: usize,
) -> Option<String> {
let qv = embedder.embed(query).await.ok()?;
let mut hits = index.search(&qv, top_k.saturating_mul(RERANK_OVERFETCH));
rerank(query, &mut hits);
hits.truncate(top_k);
render_hits(&hits, CODE_EVIDENCE_CAP)
}
pub fn gather_code_files(workspace: &str, extensions: &[String]) -> Vec<(String, String)> {
const MAX_FILES: usize = 400;
const MAX_BYTES: u64 = 200_000;
let mut out = Vec::new();
for entry in ignore::WalkBuilder::new(workspace).build().flatten() {
if out.len() >= MAX_FILES {
break;
}
let path = entry.path();
let ext = path.extension().and_then(|e| e.to_str());
if !ext.is_some_and(|e| extensions.iter().any(|x| x == e)) {
continue;
}
if path.metadata().map(|m| m.len()).unwrap_or(u64::MAX) > MAX_BYTES {
continue;
}
if let Ok(src) = std::fs::read_to_string(path) {
let rel = path
.strip_prefix(workspace)
.unwrap_or(path)
.to_string_lossy()
.to_string();
out.push((rel, src));
}
}
out
}
#[derive(Clone, Copy)]
pub struct CodeSearch<'a> {
pub embedder: &'a dyn Embedder,
pub index: &'a dyn SemanticIndex,
pub top_k: usize,
}
pub fn code_search_tool_definition() -> serde_json::Value {
serde_json::json!({
"type": "function",
"function": {
"name": "code_search",
"description": "Search the indexed codebase for the most relevant code by MEANING \
(semantic/embedding search, not keyword) — use it to find where \
something is implemented when you don't have a file path, e.g. \
'where is the retry backoff computed'. Returns the top matching \
code chunks with their file:line; then read_file the ones you need.",
"parameters": {
"type": "object",
"properties": {
"query": { "type": "string", "description": "What to find, in natural language or a symbol name." }
},
"required": ["query"]
}
}
})
}
pub(crate) async fn execute_code_search(
args: &serde_json::Value,
search: CodeSearch<'_>,
color: bool,
tool_output_lines: usize,
) -> String {
let query = args["query"].as_str().unwrap_or("").trim();
print_tool_call("code_search", query, color);
if query.is_empty() {
return "error: code_search requires a non-empty `query`".to_string();
}
let out = match retrieve_evidence(query, search.embedder, search.index, search.top_k).await {
Some(block) => block,
None => "no code matched — the semantic index may be empty or the embedding model \
unavailable; use read_file/find if you already know the path"
.to_string(),
};
print_tool_output(&out, tool_output_lines, color);
out
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use wiremock::matchers::{method, path};
use wiremock::{Mock, MockServer, Request, Respond, ResponseTemplate};
#[test]
fn gather_code_files_honors_the_extension_allowlist_956() {
static N: AtomicUsize = AtomicUsize::new(0);
let dir = std::env::temp_dir().join(format!(
"newt-gcf-{}-{}",
std::process::id(),
N.fetch_add(1, Ordering::Relaxed)
));
std::fs::create_dir_all(&dir).unwrap();
std::fs::write(dir.join("tool.sh"), "myfunc() { echo hi; }\n").unwrap();
std::fs::write(dir.join("main.rs"), "pub fn open() {}\n").unwrap();
let ws = dir.to_string_lossy().to_string();
let sh_only = gather_code_files(&ws, &["sh".to_string()]);
assert!(
sh_only.iter().any(|(p, _)| p.ends_with("tool.sh")),
"the .sh file must be gathered when `sh` is requested: {sh_only:?}"
);
assert!(
!sh_only.iter().any(|(p, _)| p.ends_with("main.rs")),
"rs was not requested: {sh_only:?}"
);
let both = gather_code_files(&ws, &["rs".to_string(), "sh".to_string()]);
assert!(both.iter().any(|(p, _)| p.ends_with("tool.sh")));
assert!(both.iter().any(|(p, _)| p.ends_with("main.rs")));
assert!(gather_code_files(&ws, &[]).is_empty());
let _ = std::fs::remove_dir_all(&dir);
}
#[tokio::test]
async fn embed_parses_the_vector() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/embeddings"))
.respond_with(
ResponseTemplate::new(200)
.set_body_json(serde_json::json!({ "embedding": [0.1, 0.2, 0.3] })),
)
.mount(&server)
.await;
let c = EmbeddingsClient::new(
server.uri(),
"nomic-embed-text",
crate::BackendKind::Ollama,
None,
30,
0,
);
assert_eq!(c.embed("hello").await.unwrap(), vec![0.1f32, 0.2, 0.3]);
}
#[tokio::test]
async fn embed_openai_protocol_hits_v1_and_parses_data() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/v1/embeddings"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"data": [{ "embedding": [0.5, 0.6] }]
})))
.mount(&server)
.await;
let c = EmbeddingsClient::new(
server.uri(),
"bge-m3",
crate::BackendKind::Openai,
None,
30,
0,
);
assert_eq!(c.embed("hello").await.unwrap(), vec![0.5f32, 0.6]);
let reqs = server.received_requests().await.unwrap();
let body: serde_json::Value = serde_json::from_slice(&reqs[0].body).unwrap();
assert_eq!(body["input"], "hello");
assert!(body.get("prompt").is_none());
}
#[tokio::test]
async fn embed_batch_preserves_order() {
struct ByLen;
impl Respond for ByLen {
fn respond(&self, req: &Request) -> ResponseTemplate {
let body: serde_json::Value = serde_json::from_slice(&req.body).unwrap_or_default();
let n = body["prompt"].as_str().unwrap_or("").chars().count() as f64;
ResponseTemplate::new(200).set_body_json(serde_json::json!({ "embedding": [n] }))
}
}
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/embeddings"))
.respond_with(ByLen)
.mount(&server)
.await;
let c = EmbeddingsClient::new(server.uri(), "m", crate::BackendKind::Ollama, None, 30, 0);
let out = c
.embed_batch(&["a".into(), "bbb".into(), "cc".into()])
.await
.unwrap();
assert_eq!(out, vec![vec![1.0f32], vec![3.0], vec![2.0]]);
}
#[tokio::test]
async fn embed_retries_then_succeeds() {
struct FailOnce(Arc<AtomicUsize>);
impl Respond for FailOnce {
fn respond(&self, _req: &Request) -> ResponseTemplate {
if self.0.fetch_add(1, Ordering::SeqCst) == 0 {
ResponseTemplate::new(500)
} else {
ResponseTemplate::new(200)
.set_body_json(serde_json::json!({ "embedding": [1.0, 2.0] }))
}
}
}
let calls = Arc::new(AtomicUsize::new(0));
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/embeddings"))
.respond_with(FailOnce(calls.clone()))
.mount(&server)
.await;
let c = EmbeddingsClient::new(server.uri(), "m", crate::BackendKind::Ollama, None, 30, 2);
assert_eq!(c.embed("x").await.unwrap(), vec![1.0f32, 2.0]);
assert_eq!(calls.load(Ordering::SeqCst), 2, "one failure, one success");
}
#[tokio::test]
async fn embed_gives_up_after_retries() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/embeddings"))
.respond_with(ResponseTemplate::new(500))
.mount(&server)
.await;
let c = EmbeddingsClient::new(server.uri(), "m", crate::BackendKind::Ollama, None, 30, 1);
let err = c.embed("x").await.unwrap_err();
assert!(err.to_string().contains("returned 500"), "{err}");
}
#[tokio::test]
async fn embed_rejects_missing_and_empty_vector() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/embeddings"))
.respond_with(
ResponseTemplate::new(200).set_body_json(serde_json::json!({ "nope": 1 })),
)
.mount(&server)
.await;
let c = EmbeddingsClient::new(server.uri(), "m", crate::BackendKind::Ollama, None, 30, 0);
assert!(c
.embed("x")
.await
.unwrap_err()
.to_string()
.contains("missing `embedding`"));
}
#[tokio::test]
async fn embed_retries_zero_makes_exactly_one_call() {
struct Count(Arc<AtomicUsize>);
impl Respond for Count {
fn respond(&self, _req: &Request) -> ResponseTemplate {
self.0.fetch_add(1, Ordering::SeqCst);
ResponseTemplate::new(500)
}
}
let calls = Arc::new(AtomicUsize::new(0));
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/embeddings"))
.respond_with(Count(calls.clone()))
.mount(&server)
.await;
let c = EmbeddingsClient::new(server.uri(), "m", crate::BackendKind::Ollama, None, 30, 0);
assert!(c.embed("x").await.is_err());
assert_eq!(
calls.load(Ordering::SeqCst),
1,
"retries=0 → exactly one attempt"
);
}
#[tokio::test]
async fn embed_sends_bearer_when_api_key_set() {
use wiremock::matchers::header;
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/embeddings"))
.and(header("authorization", "Bearer sk-test"))
.respond_with(
ResponseTemplate::new(200).set_body_json(serde_json::json!({ "embedding": [1.0] })),
)
.mount(&server)
.await;
let c = EmbeddingsClient::new(
server.uri(),
"m",
crate::BackendKind::Ollama,
Some("sk-test".into()),
30,
0,
);
assert_eq!(c.embed("x").await.unwrap(), vec![1.0f32]);
}
#[test]
fn chunk_rust_captures_defs_with_leading_doc() {
let src = "\
//! file header
use std::fmt;
/// Adds two numbers.
fn add(a: i32, b: i32) -> i32 {
a + b
}
/// A point.
struct Point {
x: i32,
}
";
let chunks = chunk_source("src/lib.rs", src);
assert_eq!(chunks.len(), 2, "one chunk per def: {chunks:#?}");
assert_eq!(chunks[0].kind, "function");
assert_eq!(chunks[0].start_line, 4);
assert!(chunks[0].text.contains("/// Adds two numbers."));
assert!(chunks[0].text.contains("fn add"));
assert!(
!chunks[0].text.contains("struct Point"),
"def boundary respected"
);
assert_eq!(chunks[1].kind, "struct");
assert!(chunks[1].text.contains("/// A point.") && chunks[1].text.contains("struct Point"));
}
#[test]
fn chunk_python_def_and_class() {
let src = "\
import os
def greet(name):
return f\"hi {name}\"
class Dog:
def bark(self):
return \"woof\"
";
let chunks = chunk_source("app.py", src);
assert!(chunks
.iter()
.any(|c| c.kind == "function" && c.text.contains("def greet")));
assert!(chunks
.iter()
.any(|c| c.kind == "class" && c.text.contains("class Dog")));
}
#[test]
fn chunk_unknown_language_falls_back_to_windows() {
let src = (1..=90)
.map(|i| format!("line {i}"))
.collect::<Vec<_>>()
.join("\n");
let chunks = chunk_source("notes.txt", &src); assert!(chunks.iter().all(|c| c.kind == "window"));
assert!(
chunks.len() >= 2,
"90 lines / 40 ⇒ 3 windows: {}",
chunks.len()
);
assert_eq!(chunks.first().unwrap().start_line, 1);
assert_eq!(chunks.last().unwrap().end_line, 90);
}
#[test]
fn chunk_oversized_body_splits_into_windows() {
let body = (0..400)
.map(|i| format!(" let v{i} = {i};"))
.collect::<Vec<_>>()
.join("\n");
let src = format!("fn big() {{\n{body}\n}}\n");
let chunks = chunk_source("src/big.rs", &src);
assert!(chunks.len() > 1, "oversized body split: {}", chunks.len());
assert!(
chunks
.iter()
.all(|c| c.text.chars().count() <= CHUNK_MAX_CHARS + 200),
"every chunk stays bounded"
);
}
#[test]
fn chunk_empty_source_is_empty() {
assert!(chunk_source("src/lib.rs", "").is_empty());
}
fn chunk(file: &str, text: &str) -> CodeChunk {
CodeChunk {
file: file.into(),
start_line: 1,
end_line: 1,
kind: "function".into(),
text: text.into(),
}
}
#[test]
fn cosine_known_vectors() {
assert!(
(cosine(&[1.0, 0.0], &[1.0, 0.0]) - 1.0).abs() < 1e-6,
"identical"
);
assert!(
cosine(&[1.0, 0.0], &[0.0, 1.0]).abs() < 1e-6,
"orthogonal → 0"
);
assert!(
(cosine(&[1.0, 0.0], &[-1.0, 0.0]) + 1.0).abs() < 1e-6,
"opposite → -1"
);
assert_eq!(cosine(&[1.0, 2.0, 3.0], &[1.0, 2.0]), 0.0);
assert_eq!(cosine(&[], &[]), 0.0);
assert_eq!(cosine(&[0.0, 0.0], &[1.0, 1.0]), 0.0);
}
#[test]
fn index_search_orders_by_cosine_and_truncates() {
let idx = SessionSemanticIndex::default();
idx.index_chunk(chunk("a.rs", "alpha"), vec![1.0, 0.0]);
idx.index_chunk(chunk("b.rs", "beta"), vec![0.0, 1.0]);
idx.index_chunk(chunk("c.rs", "gamma"), vec![0.9, 0.1]);
assert_eq!(idx.chunks_indexed(), 3);
assert_eq!(idx.indexed_chars(), 5 + 4 + 5);
let hits = idx.search(&[1.0, 0.0], 2);
assert_eq!(hits.len(), 2, "top_k truncates");
assert_eq!(hits[0].1.file, "a.rs");
assert_eq!(hits[1].1.file, "c.rs");
assert!(hits[0].0 > hits[1].0, "descending score");
assert_eq!(idx.search(&[1.0, 0.0], 99).len(), 3);
let empty = SessionSemanticIndex::default();
assert!(empty.search(&[1.0, 0.0], 5).is_empty());
idx.clear();
assert_eq!(idx.chunks_indexed(), 0);
}
#[test]
fn code_evidence_block_none_when_empty_renders_and_caps() {
let idx = SessionSemanticIndex::default();
assert_eq!(build_code_evidence_block(&idx, &[1.0, 0.0], 5, 6_000), None);
idx.index_chunk(chunk("src/lib.rs", "fn add() {}"), vec![1.0, 0.0]);
let block = build_code_evidence_block(&idx, &[1.0, 0.0], 5, 6_000).unwrap();
assert!(block.starts_with("<code_evidence>\n") && block.ends_with("</code_evidence>"));
assert!(
block.contains("src/lib.rs:1-1"),
"file:line header: {block}"
);
assert!(block.contains("fn add() {}"));
let capped = build_code_evidence_block(&idx, &[1.0, 0.0], 5, 30).unwrap();
assert!(capped.contains("omitted to fit"), "{capped}");
}
struct MockEmbedder;
#[async_trait]
impl Embedder for MockEmbedder {
async fn embed(&self, text: &str) -> anyhow::Result<Vec<f32>> {
Ok(vec![
text.matches('a').count() as f32,
text.matches('b').count() as f32,
text.chars().count() as f32,
])
}
}
struct FailEmbedder;
#[async_trait]
impl Embedder for FailEmbedder {
async fn embed(&self, _text: &str) -> anyhow::Result<Vec<f32>> {
anyhow::bail!("embeddings model not available")
}
}
struct CountingFailEmbedder(Arc<AtomicUsize>);
#[async_trait]
impl Embedder for CountingFailEmbedder {
async fn embed(&self, _text: &str) -> anyhow::Result<Vec<f32>> {
self.0.fetch_add(1, Ordering::SeqCst);
anyhow::bail!("embeddings endpoint returned 404")
}
}
#[tokio::test]
async fn index_files_disable_stops_on_first_failure() {
let files = vec![("a.rs".to_string(), "fn add() {}\nfn sub() {}".to_string())];
let idx = SessionSemanticIndex::default();
let calls = Arc::new(AtomicUsize::new(0));
let n = index_files(
&files,
&CountingFailEmbedder(calls.clone()),
&idx,
crate::OnEmbedFailure::Disable,
)
.await;
assert_eq!(n, 0);
assert_eq!(idx.chunks_indexed(), 0);
assert_eq!(
calls.load(Ordering::SeqCst),
1,
"Disable must short-circuit after the first failure"
);
let calls2 = Arc::new(AtomicUsize::new(0));
index_files(
&files,
&CountingFailEmbedder(calls2.clone()),
&SessionSemanticIndex::default(),
crate::OnEmbedFailure::Warn,
)
.await;
assert!(
calls2.load(Ordering::SeqCst) >= 2,
"Warn keeps trying per-chunk, got {}",
calls2.load(Ordering::SeqCst)
);
}
#[tokio::test]
async fn index_files_chunks_embeds_and_skips_failures() {
let files = vec![("a.rs".to_string(), "fn add() {}\nfn sub() {}".to_string())];
let idx = SessionSemanticIndex::default();
let n = index_files(&files, &MockEmbedder, &idx, crate::OnEmbedFailure::Disable).await;
assert_eq!(n, idx.chunks_indexed() as usize);
assert!(n >= 2, "two fns indexed, got {n}");
let empty = SessionSemanticIndex::default();
assert_eq!(
index_files(&files, &FailEmbedder, &empty, crate::OnEmbedFailure::Warn).await,
0
);
assert_eq!(empty.chunks_indexed(), 0);
}
#[tokio::test]
async fn retrieve_evidence_embeds_query_ranks_and_degrades() {
let files = vec![("a.rs".to_string(), "fn aaa() {}\nfn bbb() {}".to_string())];
let idx = SessionSemanticIndex::default();
index_files(&files, &MockEmbedder, &idx, crate::OnEmbedFailure::Disable).await;
let block = retrieve_evidence("aaaaa", &MockEmbedder, &idx, 1)
.await
.unwrap();
assert!(
block.contains("<code_evidence>") && block.contains("aaa"),
"{block}"
);
assert!(retrieve_evidence("x", &FailEmbedder, &idx, 1)
.await
.is_none());
let empty = SessionSemanticIndex::default();
assert!(retrieve_evidence("aaa", &MockEmbedder, &empty, 1)
.await
.is_none());
}
#[tokio::test]
async fn code_search_tool_embeds_searches_and_coaches() {
let files = vec![("a.rs".to_string(), "fn aaa() {}\nfn bbb() {}".to_string())];
let idx = SessionSemanticIndex::default();
index_files(&files, &MockEmbedder, &idx, crate::OnEmbedFailure::Disable).await;
let search = CodeSearch {
embedder: &MockEmbedder,
index: &idx,
top_k: 1,
};
let out =
execute_code_search(&serde_json::json!({"query": "aaaaa"}), search, false, 20).await;
assert!(
out.contains("<code_evidence>") && out.contains("aaa"),
"{out}"
);
assert!(
execute_code_search(&serde_json::json!({}), search, false, 20)
.await
.starts_with("error:")
);
let empty = SessionSemanticIndex::default();
let s2 = CodeSearch {
embedder: &MockEmbedder,
index: &empty,
top_k: 1,
};
assert!(
execute_code_search(&serde_json::json!({"query": "x"}), s2, false, 20)
.await
.contains("no code matched")
);
}
#[test]
fn code_search_tool_definition_shape() {
let def = code_search_tool_definition();
assert_eq!(def["function"]["name"], "code_search");
assert!(def["function"]["parameters"]["properties"]["query"].is_object());
}
#[test]
fn rerank_boosts_defs_and_paths_and_stays_stable() {
let c = |file: &str, kind: &str| CodeChunk {
file: file.to_string(),
start_line: 1,
end_line: 2,
kind: kind.to_string(),
text: "x".to_string(),
};
let mut hits = vec![(0.5, c("a.rs", "window")), (0.5, c("b.rs", "function"))];
rerank("anything", &mut hits);
assert_eq!(hits[0].1.kind, "function", "def promoted over window");
let mut hits = vec![
(0.50, c("other.rs", "window")),
(0.48, c("retry.rs", "window")),
];
rerank("where is retry handled", &mut hits);
assert_eq!(hits[0].1.file, "retry.rs", "path-term match promoted");
let mut hits = vec![(0.90, c("x.rs", "window")), (0.50, c("y.rs", "function"))];
rerank("anything", &mut hits);
assert_eq!(
hits[0].1.file, "x.rs",
"strong cosine wins over a small boost"
);
let mut hits = vec![
(0.9, c("a.rs", "window")),
(0.6, c("b.rs", "window")),
(0.4, c("c.rs", "window")),
];
let before = hits.clone();
rerank("zz", &mut hits); assert_eq!(hits, before, "no boost → cosine order unchanged");
let mut hits = vec![
(0.5, c("first.rs", "window")),
(0.5, c("second.rs", "window")),
];
rerank("zz", &mut hits);
assert_eq!(hits[0].1.file, "first.rs", "stable: ties keep input order");
}
}