feat: add Google code retrieval across desktop CLI and MCP

crates/searchi-semantic/native/worker.cpp

@@ -197,11 +197,50 @@ public:
         require(!text.empty() && text.size() <= max_passage_bytes, "Passage exceeds 1 MiB or is empty");
         auto tokens = tokenize(text,false);
         require(!tokens.empty(), "Passage has no tokens");
-        // Preserve actual source offsets, rather than inventing locations from
-        // a lossy decoded representation. Unexpected normalization is an error.
-        require(decode(tokens,tokens.size()) == text, "Tokenizer roundtrip changed source text; cannot verify source offsets");
         size_t prompt = tokenize(document_prefix(title),true).size();
         require(prompt + 132 < size_t(context_tokens), "Document prefix exceeds token budget");
+        // SentencePiece may normalize source characters. In that case decoded
+        // token lengths are not source offsets: split the original UTF-8 bytes
+        // and verify each complete prefixed fragment against the real tokenizer.
+        if (decode(tokens,tokens.size()) != text) {
+            const auto prefix = document_prefix(title);
+            json result = json::array();
+            if (prefix.size()+text.size() <= max_text_bytes && tokenize(prefix+text,true).size() <= size_t(context_tokens)) {
+                return json::array({{{"text",text},{"start_byte",0},{"end_byte",text.size()}}});
+            }
+            std::vector<size_t> boundaries;
+            for (size_t i=0;i<text.size();++i) if ((uint8_t(text[i]) & 0xc0) != 0x80) boundaries.push_back(i);
+            boundaries.push_back(text.size());
+            size_t start_index = 0;
+            while (start_index+1 < boundaries.size()) {
+                const size_t start = boundaries[start_index];
+                require(result.size() < max_windows, "Passage exceeds 256 windows");
+                size_t low = start_index;
+                size_t high = size_t(std::upper_bound(boundaries.begin(),boundaries.end(),std::min(text.size(),start+max_text_bytes-prefix.size()))-boundaries.begin())-1;
+                size_t best_index = start_index;
+                while (low < high) {
+                    size_t candidate_index = low + (high-low+1)/2;
+                    size_t candidate = boundaries[candidate_index];
+                    if (tokenize(prefix+text.substr(start,candidate-start),true).size() <= size_t(context_tokens)-4) {
+                        best_index = candidate_index;
+                        low = candidate_index;
+                    } else {
+                        high = candidate_index-1;
+                    }
+                }
+                size_t best = boundaries[best_index];
+                require(best > start, "Cannot fit source fragment in model context");
+                auto fragment = text.substr(start,best-start);
+                require(tokenize(prefix+fragment,true).size() <= size_t(context_tokens), "Source fragment exceeds context budget");
+                result.push_back({{"text",fragment},{"start_byte",start},{"end_byte",best}});
+                if (best == text.size()) break;
+                // Bounded original-byte overlap; always advance, even for
+                // unusual Unicode/token expansion. No decoded text is indexed.
+                size_t next = best - std::min(size_t(128),(best-start)/2);
+                start_index = size_t(std::lower_bound(boundaries.begin(),boundaries.end(),next)-boundaries.begin());
+            }
+            return result;
+        }
         size_t budget = size_t(context_tokens) - prompt - 4;
         size_t stride = budget - 128;
         json result = json::array();
@@ -253,10 +292,11 @@ public:
             return {{"vectors",json::array({embed_bitmap(bitmap.get())})}};
         }
         if (op == "windows") return {{"windows",windows(request.at("text").get<std::string>(),request.value("title",std::string()))}};
-        if (op == "query") {
+        if (op == "query" || op == "code_query") {
             bool debug = request.value("debug_performance",false);
             json timings;
-            auto vector = embed("task: search result | query: "+request.at("text").get<std::string>(),512,debug ? &timings : nullptr);
+            auto prefix = op == "code_query" ? "task: code retrieval | query: " : "task: search result | query: ";
+            auto vector = embed(prefix+request.at("text").get<std::string>(),512,debug ? &timings : nullptr);
             json response = {{"vectors",json::array({vector})}};
             if (debug) response["performance"] = timings;
             return response;
@@ -289,7 +329,7 @@ int main(int argc, char ** argv) {
         bool force_cpu = std::getenv("SEARCHI_FORCE_CPU") && std::string(std::getenv("SEARCHI_FORCE_CPU")) == "1";
         Encoder encoder(path,threads,device == "auto" && !force_cpu);
         std::cout << json({{"ready",true},{"model","google/embeddinggemma-2"},{"dimensions",dimensions},
-            {"device",encoder.device},{"threads",threads},{"runtime",SEARCHI_LLAMA_TAG},{"media_supported",true}}).dump() << '\n' << std::flush;
+            {"device",encoder.device},{"threads",threads},{"runtime",SEARCHI_LLAMA_TAG},{"media_supported",true},{"code_query_supported",true}}).dump() << '\n' << std::flush;
         for (;;) {
             std::string line; char c;
             while (std::cin.get(c) && c != '\n') {

crates/searchi-semantic/src/gemma.rs

@@ -48,6 +48,13 @@ pub struct WorkerInfo {
     pub runtime: String,
     #[serde(default)]
     pub media_supported: bool,
+    #[serde(default)]
+    pub code_query_supported: bool,
+}
+
+fn query_operation(info: &WorkerInfo, code: bool) -> Result<&'static str> {
+    ensure!(!code || info.code_query_supported,"Installed embedding worker does not support code retrieval; repair/rebuild the native runtime");
+    Ok(if code { "code_query" } else { "query" })
 }
 
 /// Persistent isolated native worker. No HTTP or remote model fallback.
@@ -317,15 +324,24 @@ impl GemmaEmbedder {
 
     pub fn embed_query(&mut self, query: &str) -> Result<Vec<f32>> {
         self.query_profile = None;
-        self.embed_query_inner(query, false)
+        self.embed_query_inner(query, false, false)
     }
 
     /// Opt-in diagnostics contain only durations and counts, never query text.
     pub fn embed_query_debug(&mut self, query: &str) -> Result<Vec<f32>> {
         self.query_profile = Some(BTreeMap::new());
-        self.embed_query_inner(query, true)
+        self.embed_query_inner(query, true, false)
     }
 
+    /// Google code retrieval task; unavailable workers fail explicitly.
+    pub fn embed_code_query(&mut self, query: &str) -> Result<Vec<f32>> {
+        self.query_profile = None;
+        self.embed_query_inner(query, false, true)
+    }
+    pub fn embed_code_query_debug(&mut self, query: &str) -> Result<Vec<f32>> {
+        self.query_profile = Some(BTreeMap::new());
+        self.embed_query_inner(query, true, true)
+    }
     pub fn query_diagnostics(&self) -> BTreeMap<String, u64> {
         self.query_profile.clone().unwrap_or_default()
     }
@@ -337,13 +353,13 @@ impl GemmaEmbedder {
         }
     }
 
-    fn embed_query_inner(&mut self, query: &str, debug: bool) -> Result<Vec<f32>> {
+    fn embed_query_inner(&mut self, query: &str, debug: bool, code: bool) -> Result<Vec<f32>> {
         ensure!(
             !query.trim().is_empty() && query.len() <= 4096,
             "Query must be nonempty and at most 4096 UTF-8 bytes"
         );
-        let response =
-            self.request(json!({"op":"query","text":query,"debug_performance":debug}))?;
+        let op = query_operation(&self.info, code)?;
+        let response = self.request(json!({"op":op,"text":query,"debug_performance":debug}))?;
         if debug {
             if let Some(profile) = &mut self.query_profile {
                 for name in [
@@ -655,6 +671,12 @@ mod tests {
         let mut handshake = json!({"ready":true,"model":"google/embeddinggemma-2","dimensions":768,"device":"cpu","threads":1,"runtime":"b11461-no-logits-v1"});
         let old: WorkerInfo = serde_json::from_value(handshake.clone()).unwrap();
         assert!(!old.media_supported);
+        assert!(!old.code_query_supported);
+        assert!(query_operation(&old, true).is_err());
+        assert_eq!(query_operation(&old, false).unwrap(), "query");
+        let mut new = old.clone();
+        new.code_query_supported = true;
+        assert_eq!(query_operation(&new, true).unwrap(), "code_query");
         handshake["media_supported"] = json!(true);
         assert!(
             serde_json::from_value::<WorkerInfo>(handshake)

crates/searchi-semantic/tests/code_cpu.rs

@@ -0,0 +1,66 @@
+use anyhow::Result;
+use searchi_semantic::{Device, GemmaEmbedder};
+
+#[test]
+#[ignore = "Requires explicitly provisioned CPU Google tokenizer and worker"]
+fn normalized_code_tokenization_preserves_original_source_windows() -> Result<()> {
+    let models = std::env::var_os("SEARCHI_GEMMA_MODEL_DIR").expect("Set model directory");
+    let worker = std::env::var_os("SEARCHI_GEMMA_WORKER").expect("Set CPU worker");
+    let mut embedder = GemmaEmbedder::load_with_device(models, worker, 1, Device::Cpu)?;
+    // Genuine frozen TypeScript fixture that does not roundtrip through the
+    // tokenizer. Do not replace source bytes with normalized decoded tokens.
+    let code = "export function process(rows: {category: string; amount: number}[]) {\n  const totals = new Map<string, number>();\n  for (const row of rows) totals.set(row.category, (totals.get(row.category) ?? 0) + row.amount);\n  return totals;\n}\n";
+    for text in [
+        code.to_owned(),
+        format!("{code}// Grüße 日本語 🦀\n").repeat(160),
+    ] {
+        let windows = embedder.split_windows(&text, "c.ts")?;
+        let mut covered = 0;
+        for window in &windows {
+            assert!(window.start_byte <= covered);
+            assert!(window.end_byte > covered);
+            assert_eq!(window.text, text[window.start_byte..window.end_byte]);
+            covered = window.end_byte;
+        }
+        assert_eq!(covered, text.len());
+        assert!(windows.len() <= 256);
+        if text.len() > 1000 {
+            assert!(windows.len() > 1);
+        }
+    }
+    Ok(())
+}
+
+#[test]
+#[ignore = "Requires explicitly provisioned CPU Google model and rebuilt code-capable worker"]
+fn actual_code_task_differs_from_search_and_debug_preserves_vectors() -> Result<()> {
+    let models = std::env::var_os("SEARCHI_GEMMA_MODEL_DIR").expect("Set model directory");
+    let worker = std::env::var_os("SEARCHI_GEMMA_WORKER").expect("Set CPU worker");
+    let mut embedder = GemmaEmbedder::load_with_device(models, worker, 1, Device::Cpu)?;
+    assert!(embedder.info.code_query_supported);
+    assert_eq!(embedder.info.device, "cpu");
+    let query = "parse a JSON response and retry a failed network request";
+    let ordinary = embedder.embed_query(query)?;
+    let code = embedder.embed_code_query(query)?;
+    let cosine = |a: &[f32], b: &[f32]| {
+        a.iter()
+            .zip(b)
+            .map(|(a, b)| f64::from(*a) * f64::from(*b))
+            .sum::<f64>()
+    };
+    assert!(
+        cosine(&ordinary, &code) < 0.999999,
+        "Task-specific prompt must not silently reuse generic search embeddings"
+    );
+    let debug = embedder.embed_code_query_debug(query)?;
+    assert!(cosine(&code, &debug) > 0.999999);
+    assert!(embedder
+        .query_diagnostics()
+        .contains_key("worker_token_count"));
+    let ordinary_again = embedder.embed_query(query)?;
+    assert!(
+        cosine(&ordinary, &ordinary_again) > 0.999999,
+        "Previous code query must not contaminate a retained encoder context"
+    );
+    Ok(())
+}