use std::fs;
use std::process::Command;
use std::sync::Arc;
use lantern::embed::{EmbedOptions, MockBackendFactory, MockEmbeddingBackend, embed_missing_with};
use lantern::ingest::ingest_path;
use lantern::mcp::{LanternServer, SearchArgs};
use lantern::search::{SearchOptions, search};
use lantern::store::Store;
use rusqlite::params;
use tempfile::tempdir;
fn seed_source(store: &mut Store, source_id: &str) {
let conn = store.conn();
conn.execute(
"INSERT INTO sources (id, uri, path, kind, bytes, content_sha256, mtime_unix, ingested_at)
VALUES (?1, ?2, NULL, 'application/jsonl', 0, ?3, NULL, 0)",
params![
source_id,
format!("mem://{source_id}"),
format!("sha-{source_id}")
],
)
.unwrap();
}
fn seed_chunk(
store: &mut Store,
chunk_id: &str,
source_id: &str,
ordinal: i64,
session_id: Option<&str>,
text: &str,
) {
let byte_end = text.len() as i64;
let conn = store.conn();
conn.execute(
"INSERT INTO chunks (id, source_id, ordinal, byte_start, byte_end, char_count, text,
sha256, created_at, session_id)
VALUES (?1, ?2, ?3, 0, ?4, ?5, ?6, ?7, 0, ?8)",
params![
chunk_id,
source_id,
ordinal,
byte_end,
text.chars().count() as i64,
text,
format!("chunk-sha-{chunk_id}"),
session_id,
],
)
.unwrap();
}
#[test]
fn search_session_id_filter_restricts_to_single_session() {
let dir = tempdir().unwrap();
let mut store = Store::initialize(&dir.path().join("store")).unwrap();
seed_source(&mut store, "src1");
seed_chunk(&mut store, "c1", "src1", 0, Some("sess-A"), "needle alpha");
seed_chunk(&mut store, "c2", "src1", 1, Some("sess-B"), "needle beta");
seed_chunk(&mut store, "c3", "src1", 2, None, "needle uncovered");
let opts = SearchOptions {
session_id: Some("sess-A".into()),
..SearchOptions::default()
};
let hits = search(&store, "needle", opts).unwrap();
assert_eq!(
hits.len(),
1,
"session filter should keep one chunk; got {hits:#?}"
);
assert_eq!(hits[0].chunk_id, "c1");
assert_eq!(hits[0].session_id.as_deref(), Some("sess-A"));
}
#[test]
fn search_session_id_filter_excludes_chunks_without_session() {
let dir = tempdir().unwrap();
let mut store = Store::initialize(&dir.path().join("store")).unwrap();
seed_source(&mut store, "src1");
seed_chunk(&mut store, "c1", "src1", 0, None, "needle floats free");
seed_chunk(
&mut store,
"c2",
"src1",
1,
Some("sess-A"),
"needle in sess-A",
);
let opts = SearchOptions {
session_id: Some("sess-A".into()),
..SearchOptions::default()
};
let hits = search(&store, "needle", opts).unwrap();
assert_eq!(hits.len(), 1);
assert_eq!(hits[0].chunk_id, "c2");
}
#[test]
fn search_unknown_session_id_returns_empty() {
let dir = tempdir().unwrap();
let mut store = Store::initialize(&dir.path().join("store")).unwrap();
seed_source(&mut store, "src1");
seed_chunk(&mut store, "c1", "src1", 0, Some("sess-A"), "needle alpha");
let opts = SearchOptions {
session_id: Some("sess-missing".into()),
..SearchOptions::default()
};
let hits = search(&store, "needle", opts).unwrap();
assert!(hits.is_empty());
}
#[test]
fn search_default_options_ignore_session_id() {
let dir = tempdir().unwrap();
let mut store = Store::initialize(&dir.path().join("store")).unwrap();
seed_source(&mut store, "src1");
seed_chunk(&mut store, "c1", "src1", 0, Some("sess-A"), "needle alpha");
seed_chunk(&mut store, "c2", "src1", 1, Some("sess-B"), "needle beta");
let hits = search(&store, "needle", SearchOptions::default()).unwrap();
assert_eq!(
hits.len(),
2,
"default opts must not silently filter sessions"
);
}
fn lantern() -> Command {
Command::new(env!("CARGO_BIN_EXE_lantern"))
}
fn run(mut cmd: Command) -> (String, String, bool) {
let output = cmd.output().expect("failed to run lantern");
(
String::from_utf8_lossy(&output.stdout).into_owned(),
String::from_utf8_lossy(&output.stderr).into_owned(),
output.status.success(),
)
}
fn write_jsonl(path: &std::path::Path, lines: &[&str]) {
fs::create_dir_all(path.parent().unwrap()).unwrap();
let mut body = String::new();
for line in lines {
body.push_str(line);
body.push('\n');
}
fs::write(path, body).unwrap();
}
#[test]
fn search_cli_accepts_session_id_flag() {
let root = tempdir().unwrap();
let store = root.path().join("store");
let transcript = root.path().join("data/session.jsonl");
write_jsonl(
&transcript,
&[
"{\"role\":\"user\",\"session_id\":\"sess-A\",\"content\":\"needle in sess-A\"}",
"{\"role\":\"user\",\"session_id\":\"sess-B\",\"content\":\"needle in sess-B\"}",
],
);
let (_, _, ok) = run({
let mut c = lantern();
c.args(["init", "--path"]).arg(&store);
c
});
assert!(ok, "init failed");
let (_, _, ok) = run({
let mut c = lantern();
c.arg("ingest").arg(&transcript).arg("--store").arg(&store);
c
});
assert!(ok, "ingest failed");
let (stdout, _, ok) = run({
let mut c = lantern();
c.arg("search")
.arg("needle")
.arg("--store")
.arg(&store)
.args(["--session-id", "sess-A"]);
c
});
assert!(ok, "search failed");
assert!(
stdout.contains("hits: 1"),
"session-filtered search should return one hit: {stdout}"
);
assert!(
stdout.contains("session=sess-A"),
"expected session-A metadata in output: {stdout}"
);
assert!(
!stdout.contains("session=sess-B"),
"session-B chunk should be filtered out: {stdout}"
);
}
#[test]
fn query_cli_accepts_session_id_flag() {
let root = tempdir().unwrap();
let store = root.path().join("store");
let transcript = root.path().join("data/session.jsonl");
write_jsonl(
&transcript,
&[
"{\"role\":\"user\",\"session_id\":\"sess-A\",\"content\":\"needle in sess-A\"}",
"{\"role\":\"user\",\"session_id\":\"sess-B\",\"content\":\"needle in sess-B\"}",
],
);
let (_, _, ok) = run({
let mut c = lantern();
c.args(["init", "--path"]).arg(&store);
c
});
assert!(ok);
let (_, _, ok) = run({
let mut c = lantern();
c.arg("ingest").arg(&transcript).arg("--store").arg(&store);
c
});
assert!(ok);
let (stdout, _, ok) = run({
let mut c = lantern();
c.arg("query")
.arg("needle")
.arg("--store")
.arg(&store)
.args(["--session-id", "sess-B"]);
c
});
assert!(ok);
assert!(
stdout.contains("hits: 1"),
"query session filter missing: {stdout}"
);
assert!(stdout.contains("session=sess-B"));
assert!(!stdout.contains("session=sess-A"));
}
const MCP_MOCK_MODEL: &str = "mock-embed-session";
fn ingest_two_session_transcript(root: &std::path::Path) -> std::path::PathBuf {
let store_dir = root.join("store");
let mut store = Store::initialize(&store_dir).unwrap();
let data = root.join("data");
fs::create_dir_all(&data).unwrap();
let transcript = data.join("session.jsonl");
write_jsonl(
&transcript,
&[
"{\"role\":\"user\",\"session_id\":\"sess-A\",\"content\":\"needle in sess-A\"}",
"{\"role\":\"user\",\"session_id\":\"sess-B\",\"content\":\"needle in sess-B\"}",
],
);
ingest_path(&mut store, &data).unwrap();
drop(store);
store_dir
}
#[test]
fn mcp_keyword_search_accepts_session_id_filter() {
let root = tempdir().unwrap();
let store_dir = ingest_two_session_transcript(root.path());
let server = LanternServer::new(store_dir);
let resp = server
.search_sync(SearchArgs {
query: "needle".to_string(),
limit: Some(10),
kind: None,
path: None,
mode: Some("keyword".to_string()),
model: None,
ollama_url: None,
instruction: None,
min_confidence: None,
session_id: Some("sess-A".to_string()),
})
.unwrap();
let results = resp
.get("results")
.and_then(|v| v.as_array())
.expect("results array");
assert_eq!(
results.len(),
1,
"MCP keyword session filter mismatch: {resp:#?}"
);
assert_eq!(
results[0].get("session_id").and_then(|v| v.as_str()),
Some("sess-A"),
);
}
#[test]
fn mcp_semantic_search_accepts_session_id_filter() {
let root = tempdir().unwrap();
let store_dir = ingest_two_session_transcript(root.path());
let mut store = Store::open(&store_dir).unwrap();
let backend = MockEmbeddingBackend::new(64);
embed_missing_with(
&mut store,
&EmbedOptions {
model: MCP_MOCK_MODEL.to_string(),
ollama_url: "http://mock".to_string(),
limit: None,
},
&backend,
)
.unwrap();
drop(store);
let factory = Arc::new(MockBackendFactory::new(64));
let server = LanternServer::with_factory(store_dir, factory);
let resp = server
.search_sync(SearchArgs {
query: "needle".to_string(),
limit: Some(10),
kind: None,
path: None,
mode: Some("semantic".to_string()),
model: Some(MCP_MOCK_MODEL.to_string()),
ollama_url: Some("http://mock".to_string()),
instruction: None,
min_confidence: None,
session_id: Some("sess-B".to_string()),
})
.unwrap();
let results = resp
.get("results")
.and_then(|v| v.as_array())
.expect("results array");
assert_eq!(
results.len(),
1,
"MCP semantic session filter mismatch: {resp:#?}"
);
assert_eq!(
results[0].get("session_id").and_then(|v| v.as_str()),
Some("sess-B"),
);
}
#[test]
fn mcp_hybrid_search_accepts_session_id_filter() {
let root = tempdir().unwrap();
let store_dir = ingest_two_session_transcript(root.path());
let mut store = Store::open(&store_dir).unwrap();
let backend = MockEmbeddingBackend::new(64);
embed_missing_with(
&mut store,
&EmbedOptions {
model: MCP_MOCK_MODEL.to_string(),
ollama_url: "http://mock".to_string(),
limit: None,
},
&backend,
)
.unwrap();
drop(store);
let factory = Arc::new(MockBackendFactory::new(64));
let server = LanternServer::with_factory(store_dir, factory);
let resp = server
.search_sync(SearchArgs {
query: "needle".to_string(),
limit: Some(10),
kind: None,
path: None,
mode: Some("hybrid".to_string()),
model: Some(MCP_MOCK_MODEL.to_string()),
ollama_url: Some("http://mock".to_string()),
instruction: None,
min_confidence: None,
session_id: Some("sess-A".to_string()),
})
.unwrap();
let results = resp
.get("results")
.and_then(|v| v.as_array())
.expect("results array");
assert_eq!(
results.len(),
1,
"MCP hybrid session filter mismatch: {resp:#?}"
);
assert_eq!(
results[0].get("session_id").and_then(|v| v.as_str()),
Some("sess-A"),
);
}