Use measured query-only context across desktop CLI and MCP

crates/searchi-cli/tests/mcp.rs

@@ -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"));
+}