Use measured query-only context across desktop CLI and MCP

This commit is contained in:
karl committed 2026-10-13 07:30:40 +02:00
1 parent 95c72c5fd4
commit 88cd9c7f5a
12 files changed
+582 -20

No files matched your search

+1 -1
View File
@@ -203,7 +203,7 @@ fn main() -> Result<()> {
});
let results = if let (Some(model), Some(runtime)) = (model_dir, embedding_worker) {
let mut service = trace.timed("cli.session_initialize", None, |_, _| {
HybridSearch::open(&db, model, runtime, usize::from(cpu_threads))
HybridSearch::open_query(&db, model, runtime, usize::from(cpu_threads))
})?;
service
.search_traced(&query, usize::from(limit), &options, &mut trace, None)
+1 -1
View File
@@ -450,7 +450,7 @@ pub fn serve(
"No semantic generation; run explicit indexing/embedding outside MCP first"
);
let mut hybrid = HybridSearch::open(&db, models, worker, threads)?;
let mut hybrid = HybridSearch::open_query(&db, models, worker, threads)?;
hybrid.warmup()?;
+123 -3
View File
@@ -19,9 +19,22 @@ impl Drop for Server {
}
impl Server {
fn start(db: &std::path::Path) -> Self {
let mut child = Command::new(env!("CARGO_BIN_EXE_searchi"))
.args(["mcp", "--db"])
.arg(db)
Self::start_with_assets(db, None)
}
fn start_with_assets(
db: &std::path::Path,
assets: Option<(&std::ffi::OsStr, &std::ffi::OsStr)>,
) -> Self {
let mut command = Command::new(env!("CARGO_BIN_EXE_searchi"));
command.args(["mcp", "--db"]).arg(db);
if let Some((model, worker)) = assets {
command
.arg("--model-dir")
.arg(model)
.arg("--embedding-worker")
.arg(worker);
}
let mut child = command
.env("SEARCHI_FORCE_CPU", "1")
.stdin(Stdio::piped())
.stdout(Stdio::piped())
@@ -233,3 +246,110 @@ fn missing_database_never_creates_file_or_protocol_noise() {
assert!(output.stdout.is_empty());
assert!(!db.exists());
}
#[test]
#[ignore = "Requires verified Google model and current CPU-capable native worker"]
fn actual_google_cli_and_mcp_use_query_role_with_cited_results() {
let models = std::env::var_os("SEARCHI_GEMMA_MODEL_DIR").expect("Set models");
let worker = std::env::var_os("SEARCHI_GEMMA_WORKER").expect("Set worker");
let fixture = tempfile::tempdir().unwrap();
let folder = fixture.path().join("documents");
std::fs::create_dir(&folder).unwrap();
std::fs::write(
folder.join("supplier.txt"),
"The supplier agreement may be terminated with sixty days written notice.",
)
.unwrap();
std::fs::write(
folder.join("garden.txt"),
"Vegetable beds and rainwater collection in the garden.",
)
.unwrap();
std::fs::write(
folder.join("retry.rs"),
"pub fn retry_backoff(attempt: u32) -> u64 {\n 100 * 2_u64.pow(attempt)\n}\n",
)
.unwrap();
let db = fixture.path().join("index.sqlite");
searchi_core::SearchIndex::open(&db)
.unwrap()
.index_folder(&folder)
.unwrap();
let indexer = searchi_service::HybridSearch::open(&db, &models, &worker, 1).unwrap();
assert_eq!(indexer.worker_info().unwrap().role, "indexing");
drop(indexer);
let embedded = Command::new(env!("CARGO_BIN_EXE_searchi"))
.args(["embed", "--db"])
.arg(&db)
.arg("--model-dir")
.arg(&models)
.arg("--embedding-worker")
.arg(&worker)
.arg("--text-only")
.env("SEARCHI_FORCE_CPU", "1")
.output()
.unwrap();
assert!(
embedded.status.success(),
"{}",
String::from_utf8_lossy(&embedded.stderr)
);
let output = Command::new(env!("CARGO_BIN_EXE_searchi"))
.args(["search", "--db"])
.arg(&db)
.args(["cancel the supplier partnership", "--json"])
.arg("--model-dir")
.arg(&models)
.arg("--embedding-worker")
.arg(&worker)
.env("SEARCHI_FORCE_CPU", "1")
.output()
.unwrap();
assert!(
output.status.success(),
"{}",
String::from_utf8_lossy(&output.stderr)
);
let cli: Value = serde_json::from_slice(&output.stdout).unwrap();
assert_eq!(cli[0]["filename"], "supplier.txt");
let mut server = Server::start_with_assets(&db, Some((&models, &worker)));
server.request(1,"initialize",json!({"protocolVersion":"2024-11-05","capabilities":{},"clientInfo":{"name":"query-role-regression","version":"1"}}));
server.send(json!({"jsonrpc":"2.0","method":"notifications/initialized"}));
let stats = server.request(2, "tools/call", json!({"name":"stats","arguments":{}}));
let status = &stats["result"]["structuredContent"];
assert_eq!(status["worker"]["role"], "query");
assert_eq!(status["worker"]["context_tokens"], 512);
assert_eq!(status["worker"]["device"], "cpu");
assert_eq!(status["worker"]["media_supported"], true);
let found = server.request(
3,
"tools/call",
json!({"name":"search","arguments":{"query":"cancel the supplier partnership","limit":2}}),
);
let first = &found["result"]["structuredContent"]["results"][0];
assert_eq!(first["filename"], "supplier.txt");
let handle = first["source_handle"].as_str().unwrap();
let source = server.request(
4,
"tools/call",
json!({"name":"read_source","arguments":{"handle":handle}}),
);
assert!(source["result"]["structuredContent"]["text"]
.as_str()
.unwrap()
.contains("sixty days"));
let code=server.request(5,"tools/call",json!({"name":"search","arguments":{"query":"retry failed requests with exponential backoff","code_only":true}}));
let first = &code["result"]["structuredContent"]["results"][0];
assert_eq!(first["filename"], "retry.rs");
assert_eq!(first["source"]["kind"], "line");
let handle = first["source_handle"].as_str().unwrap();
let source = server.request(
6,
"tools/call",
json!({"name":"read_source","arguments":{"handle":handle}}),
);
assert!(source["result"]["structuredContent"]["text"]
.as_str()
.unwrap()
.contains("2_u64.pow"));
}