Use measured query-only context across desktop CLI and MCP

crates/searchi-semantic/native/worker.cpp

@@ -35,6 +35,7 @@ constexpr size_t max_passage_bytes = 1024 * 1024;
 constexpr size_t max_text_bytes = 65536;
 constexpr size_t max_windows = 256;
 constexpr int context_tokens = 2048;
+constexpr int query_context_tokens = 512;
 constexpr int dimensions = 768;
 
 static void require(bool condition, const char * message) {
@@ -53,6 +54,8 @@ class Encoder {
     const llama_vocab * vocab = nullptr;
     mtmd_context * media = nullptr;
     std::string model_path;
+    bool query_only;
+    int allocated_context;
     void release() {
         if (media) mtmd_free(media);
         if (batch) llama_batch_ext_free(batch);
@@ -69,7 +72,7 @@ class Encoder {
         model = llama_model_load_from_file(model_path.c_str(),mp);
         if (!model) return false;
         auto cp = llama_context_default_params();
-        cp.n_ctx = context_tokens; cp.n_batch = context_tokens; cp.n_ubatch = context_tokens;
+        cp.n_ctx = allocated_context; cp.n_batch = allocated_context; cp.n_ubatch = allocated_context;
         cp.n_seq_max = 1; cp.kv_unified = true; cp.embeddings = true;
         cp.offload_kqv = gpu; cp.op_offload = gpu;
         cp.pooling_type = LLAMA_POOLING_TYPE_MEAN;
@@ -88,7 +91,7 @@ public:
     std::string cpu_backend_module;
     bool vulkan_module_loaded = false;
     int threads;
-    Encoder(const std::string & path, int cpu_threads, bool want_gpu) : model_path(path),threads(cpu_threads) {
+    Encoder(const std::string & path, int cpu_threads, bool want_gpu, bool query_role) : model_path(path),query_only(query_role),allocated_context(query_role ? query_context_tokens : context_tokens),threads(cpu_threads) {
         llama_log_set([](ggml_log_level level, const char * text, void *) {
             if (level >= GGML_LOG_LEVEL_WARN) std::fputs(text,stderr);
         },nullptr);
@@ -331,6 +334,7 @@ public:
     }
     json handle(const json & request) {
         auto op = request.at("op").get<std::string>();
+        require(!query_only || op == "query" || op == "code_query", "Query worker only supports query and code_query operations");
         if(op=="image" || op=="audio") {
             auto encoded=request.at("data").get<std::string>();
             require(encoded.size()<=12000000 && encoded.size()%4==0,"Invalid media payload length");
@@ -363,7 +367,7 @@ public:
             bool debug = request.value("debug_performance",false);
             json timings;
             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);
+            auto vector = embed(prefix+request.at("text").get<std::string>(),query_context_tokens,debug ? &timings : nullptr);
             json response = {{"vectors",json::array({vector})}};
             if (debug) response["performance"] = timings;
             return response;
@@ -382,17 +386,19 @@ public:
 
 int main(int argc, char ** argv) {
     try {
-        std::string path, device = "auto"; int threads = 1;
+        std::string path, device = "auto", role = "indexing"; int threads = 1;
         for (int i = 1; i < argc; ++i) {
             require(i+1 < argc,"Worker options require values");
             std::string option = argv[i++], value = argv[i];
             if (option == "--model") path = value;
             else if (option == "--threads") threads = std::stoi(value);
             else if (option == "--device") device = value;
+            else if (option == "--role") role = value;
             else throw std::runtime_error("Unknown worker option");
         }
         require(!path.empty() && threads >= 1 && threads <= 8,"Model path required; threads must be 1–8");
         require(device == "cpu" || device == "auto","Device must be cpu or auto");
+        require(role == "query" || role == "indexing", "Role must be query or indexing");
         bool force_cpu = std::getenv("SEARCHI_FORCE_CPU") && std::string(std::getenv("SEARCHI_FORCE_CPU")) == "1";
         // Disable backend registration before llama initializes its registry.
         // A GPU-capable distribution can therefore be tested on CPU without
@@ -406,9 +412,9 @@ int main(int argc, char ** argv) {
             require(setenv("GGML_DISABLE_VULKAN", "1", 1) == 0, "Cannot disable Vulkan for CPU startup");
 #endif
         }
-        Encoder encoder(path,threads,device == "auto" && !force_cpu);
+        Encoder encoder(path,threads,device == "auto" && !force_cpu,role == "query");
         std::cout << json({{"ready",true},{"model","google/embeddinggemma-2"},{"dimensions",dimensions},
-            {"device",encoder.device},{"threads",threads},{"runtime",SEARCHI_LLAMA_TAG},{"media_supported",true},{"code_query_supported",true},
+            {"device",encoder.device},{"threads",threads},{"runtime",SEARCHI_LLAMA_TAG},{"media_supported",true},{"code_query_supported",true},{"role",role},{"context_tokens",role == "query" ? query_context_tokens : context_tokens},
             {"cpu_backend_module",encoder.cpu_backend_module},{"vulkan_module_loaded",encoder.vulkan_module_loaded}}).dump() << '\n' << std::flush;
         for (;;) {
             std::string line; char c;