perf(gateway): Phase 1 search_tools — drop readiness resolve, Arc cache

Warm path no longer full-space feature-resolves for enrichment; session index is Arc-shared and warmer embeds in one batch. Measured warm HogQL ~346ms → ~33ms.

Signed-off-by: crimsonsunset <jsangio1@gmail.com>

Author: crimsonsunset

crates/mcpmux-gateway/src/services/embedding_warmer.rs

@@ -181,20 +181,37 @@ impl EmbeddingWarmer {
             return Ok(());
         }
 
-        let mut records = Vec::new();
         let embed_started = Instant::now();
+        let mut ordered_hashes = Vec::with_capacity(missing_hashes.len());
+        let mut ordered_haystacks = Vec::with_capacity(missing_hashes.len());
         for content_hash in missing_hashes {
             let Some(haystack) = haystacks_by_hash.get(&content_hash).cloned() else {
                 continue;
             };
+            ordered_hashes.push(content_hash);
+            ordered_haystacks.push(haystack);
+        }
 
-            let Some(vectors) = self.embeddings.embed_documents(&[haystack], None) else {
-                continue;
-            };
-            let Some(vector) = vectors.into_iter().next() else {
-                continue;
-            };
+        let Some(vectors) = self
+            .embeddings
+            .embed_documents(&ordered_haystacks, None)
+        else {
+            info!(
+                space_id = %space_id,
+                server_id,
+                embedded = 0,
+                skipped_present,
+                missing = ordered_hashes.len(),
+                embed_ms = embed_started.elapsed().as_millis() as u64,
+                model_version = self.embeddings.model_version(),
+                model_state = ?self.embeddings.state(),
+                "[embed] warm batch done"
+            );
+            return Ok(());
+        };
 
+        let mut records = Vec::with_capacity(vectors.len());
+        for (content_hash, vector) in ordered_hashes.into_iter().zip(vectors) {
             records.push(EmbeddingRecord {
                 content_hash: content_hash.clone(),
                 model_version: self.embeddings.model_version().to_string(),

crates/mcpmux-gateway/src/services/meta_tools/meta_tool_common.rs

@@ -181,11 +181,32 @@ pub(crate) fn format_invoke_not_ready_action_with_name(
     }
 }
 
-/// Display names and pre-configured default param keys per installed server.
-pub(crate) async fn build_installed_server_meta_maps(
+/// Search-hit enrichment built from one installed-server list + pool statuses.
+///
+/// `binding_server_ids` should be the unique `server_id`s from the caller's
+/// **active** tool index (not a full `resolve_feature_sets` catalog scan).
+pub(crate) struct SearchServerEnrichment {
+    pub readiness_map: HashMap<String, &'static str>,
+    pub display_names: HashMap<String, String>,
+    pub prefilled_params_by_server: HashMap<String, Vec<String>>,
+}
+
+/// Whether the caller omitted or blanked the search query.
+pub(crate) fn is_query_empty(query: Option<&str>) -> bool {
+    query.map(str::trim).is_none_or(str::is_empty)
+}
+
+/// Build readiness + display/prefill maps without re-resolving FeatureSets.
+///
+/// One `list_for_space` for installed servers and one `get_all_statuses` pass.
+/// Extra `server_ids` (e.g. inactive widen) are labeled `bindable` unless they
+/// also appear in `binding_server_ids`.
+pub(crate) async fn build_search_server_enrichment(
     call: &MetaToolCall<'_>,
     space_id: &Uuid,
-) -> Result<(HashMap<String, String>, HashMap<String, Vec<String>>), MetaToolError> {
+    binding_server_ids: &HashSet<String>,
+    extra_server_ids: &HashSet<String>,
+) -> Result<SearchServerEnrichment, MetaToolError> {
     let installed = call
         .ctx
         .installed_server_repo
@@ -194,63 +215,31 @@ pub(crate) async fn build_installed_server_meta_maps(
         .map_err(|e| MetaToolError::Internal(e.to_string()))?;
 
     let mut display_names = HashMap::new();
-    let mut prefilled_params = HashMap::new();
+    let mut prefilled_params_by_server = HashMap::new();
+    let mut installed_by_id: HashMap<String, mcpmux_core::InstalledServer> = HashMap::new();
     for server in installed {
         display_names.insert(server.server_id.clone(), server.display_name().to_string());
         if !server.default_params.is_empty() {
             let mut keys: Vec<String> = server.default_params.keys().cloned().collect();
             keys.sort();
-            prefilled_params.insert(server.server_id.clone(), keys);
+            prefilled_params_by_server.insert(server.server_id.clone(), keys);
         }
+        installed_by_id.insert(server.server_id.clone(), server);
     }
 
-    Ok((display_names, prefilled_params))
-}
-
-/// Whether the caller omitted or blanked the search query.
-pub(crate) fn is_query_empty(query: Option<&str>) -> bool {
-    query.map(str::trim).is_none_or(str::is_empty)
-}
-
-/// Point-in-time `readiness` label per server for search hit enrichment.
-pub(crate) async fn build_server_readiness_map(
-    call: &MetaToolCall<'_>,
-    space_id: &Uuid,
-    resolved: &ResolvedFeatureSet,
-) -> Result<HashMap<String, &'static str>, MetaToolError> {
-    let binding_features = call
-        .ctx
-        .feature_service
-        .resolve_feature_sets(&space_id.to_string(), &resolved.feature_set_ids)
-        .await?;
-    let binding_servers: HashSet<String> = binding_features
-        .iter()
-        .map(|f| f.server_id.clone())
-        .collect();
-
-    let installed = call
-        .ctx
-        .installed_server_repo
-        .list_for_space(&space_id.to_string())
-        .await
-        .map_err(|e| MetaToolError::Internal(e.to_string()))?;
-    let installed_by_id: HashMap<String, mcpmux_core::InstalledServer> = installed
-        .into_iter()
-        .map(|s| (s.server_id.clone(), s))
-        .collect();
-
     let pool_statuses = call.ctx.server_manager.get_all_statuses(*space_id).await;
 
-    let server_ids: HashSet<String> = binding_servers
+    let server_ids: HashSet<String> = binding_server_ids
         .iter()
+        .chain(extra_server_ids.iter())
         .chain(installed_by_id.keys())
         .cloned()
         .collect();
 
-    let map = server_ids
+    let readiness_map = server_ids
         .into_iter()
         .map(|server_id| {
-            let in_binding = binding_servers.contains(&server_id);
+            let in_binding = binding_server_ids.contains(&server_id);
             let connection_status = pool_statuses
                 .get(&server_id)
                 .map(|(status, _, _, _)| *status)
@@ -264,7 +253,12 @@ pub(crate) async fn build_server_readiness_map(
             (server_id, readiness)
         })
         .collect();
-    Ok(map)
+
+    Ok(SearchServerEnrichment {
+        readiness_map,
+        display_names,
+        prefilled_params_by_server,
+    })
 }
 
 /// Common path for every write tool: build payload, ask broker, run the

crates/mcpmux-gateway/src/services/meta_tools/registry.rs

@@ -82,7 +82,7 @@ pub struct MetaToolContext {
     /// Per-server log tail reader (`current.log`); same source as the desktop UI.
     pub log_manager: Arc<ServerLogManager>,
     /// Per-session active tool index for `mcpmux_search_tools` (fingerprint-keyed).
-    pub search_cache: Arc<DashMap<String, (u64, ToolIndex)>>,
+    pub search_cache: Arc<DashMap<String, (u64, Arc<ToolIndex>)>>,
     /// Global embedding vectors keyed by content hash.
     pub embedding_store: Arc<DashMap<String, Vec<f32>>>,
     /// Persistent embedding repository backing `embedding_store` hydration.

crates/mcpmux-gateway/src/services/meta_tools/search_tools.rs

@@ -1,18 +1,19 @@
 //! `mcpmux_search_tools` — hybrid search and browse over the active tool index.
 
+use std::collections::HashSet;
+use std::sync::Arc;
+use std::time::Instant;
+
 use async_trait::async_trait;
 use rmcp::model::CallToolResult;
 use serde_json::{json, Value};
-use std::collections::HashSet;
-use std::time::Instant;
 use tracing::{debug, info};
 use uuid::Uuid;
 
 use crate::services::embedding::model_state_label;
 
 use super::meta_tool_common::{
-    build_installed_server_meta_maps, build_server_readiness_map, caller_resolution,
-    is_query_empty, text_result,
+    build_search_server_enrichment, caller_resolution, is_query_empty, text_result,
 };
 use super::registry::{feature_set_ids_fingerprint, MetaTool, MetaToolCall, MetaToolError};
 use super::search_tools_index::{
@@ -161,25 +162,28 @@ impl MetaTool for SearchToolsTool {
             debug!(query_id = %query_id, query, "[search] query text");
         }
 
-        let readiness_started = Instant::now();
-        let readiness_map = build_server_readiness_map(&call, &space_id, &resolved).await?;
-        let readiness_ms = readiness_started.elapsed().as_millis() as u64;
-
-        let installed_meta_started = Instant::now();
-        let (server_display_names, prefilled_params_by_server) =
-            build_installed_server_meta_maps(&call, &space_id).await?;
-        let installed_meta_ms = installed_meta_started.elapsed().as_millis() as u64;
-
         let mut index_cache_hit = false;
         let active_index_started = Instant::now();
-        let active_index = if let Some(session_id) = call.session_id {
-            if let Some(entry) = call.ctx.search_cache.get(session_id) {
-                let (cached_fp, cached_index) = entry.value();
-                if *cached_fp == fingerprint {
-                    index_cache_hit = true;
-                    cached_index.clone()
+        let active_index: Arc<crate::services::ToolIndex> =
+            if let Some(session_id) = call.session_id {
+                if let Some(entry) = call.ctx.search_cache.get(session_id) {
+                    let (cached_fp, cached_index) = entry.value();
+                    if *cached_fp == fingerprint {
+                        index_cache_hit = true;
+                        Arc::clone(cached_index)
+                    } else {
+                        drop(entry);
+                        build_and_cache_active_index(
+                            &call,
+                            &space_id,
+                            &resolved,
+                            fingerprint,
+                            session_id,
+                            query_id.as_str(),
+                        )
+                        .await?
+                    }
                 } else {
-                    drop(entry);
                     build_and_cache_active_index(
                         &call,
                         &space_id,
@@ -191,19 +195,8 @@ impl MetaTool for SearchToolsTool {
                     .await?
                 }
             } else {
-                build_and_cache_active_index(
-                    &call,
-                    &space_id,
-                    &resolved,
-                    fingerprint,
-                    session_id,
-                    query_id.as_str(),
-                )
-                .await?
-            }
-        } else {
-            build_active_index(&call, &space_id, &resolved, query_id.as_str()).await?
-        };
+                build_active_index(&call, &space_id, &resolved, query_id.as_str()).await?
+            };
         let active_index_ms = active_index_started.elapsed().as_millis() as u64;
 
         debug!(
@@ -214,13 +207,18 @@ impl MetaTool for SearchToolsTool {
             "[search] active index ready"
         );
 
-        let clone_started = Instant::now();
-        let mut index = active_index.clone();
-        let index_clone_ms = clone_started.elapsed().as_millis() as u64;
+        let binding_server_ids: HashSet<String> = active_index
+            .iter()
+            .map(|entry| entry.server_id.clone())
+            .collect();
 
         let mut inactive_tool_count = 0usize;
         let mut inactive_widen_ms = 0_u64;
+        let mut index_clone_ms = 0_u64;
+        let mut extra_server_ids: HashSet<String> = HashSet::new();
 
+        // Copy-on-write only when inactive widen mutates the haystack.
+        let mut widened_index: Option<crate::services::ToolIndex> = None;
         if include_inactive {
             debug!(
                 query_id = %query_id,
@@ -242,12 +240,16 @@ impl MetaTool for SearchToolsTool {
                 crate::services::tool_discovery::ToolDiscoveryService::build_inactive_index(
                     &inactive,
                 );
+            let clone_started = Instant::now();
+            let mut index = (*active_index).clone();
+            index_clone_ms = clone_started.elapsed().as_millis() as u64;
             let active_keys: HashSet<(String, String)> = index
                 .iter()
                 .map(|e| (e.server_id.clone(), e.feature_name.clone()))
                 .collect();
             let before_merge = index.len();
             for entry in inactive_index {
+                extra_server_ids.insert(entry.server_id.clone());
                 let key = (entry.server_id.clone(), entry.feature_name.clone());
                 if !active_keys.contains(&key) {
                     index.push(entry);
@@ -263,8 +265,29 @@ impl MetaTool for SearchToolsTool {
                 inactive_widen_ms,
                 "[search] inactive widen complete"
             );
+            widened_index = Some(index);
         }
 
+        let index: &[crate::services::ToolIndexEntry] = widened_index
+            .as_deref()
+            .unwrap_or(active_index.as_slice());
+        let merged_index_len = index.len();
+
+        let enrichment_started = Instant::now();
+        let enrichment = build_search_server_enrichment(
+            &call,
+            &space_id,
+            &binding_server_ids,
+            &extra_server_ids,
+        )
+        .await?;
+        let readiness_ms = enrichment_started.elapsed().as_millis() as u64;
+        // installed list is folded into enrichment — keep the timing field for log continuity.
+        let installed_meta_ms = 0_u64;
+        let readiness_map = enrichment.readiness_map;
+        let server_display_names = enrichment.display_names;
+        let prefilled_params_by_server = enrichment.prefilled_params_by_server;
+
         let (hydrate_ms, hydrated_missing_count) = if effective_query.is_some() {
             let hydrate =
                 hydrate_active_embeddings(&call, query_id.as_str(), active_index.as_slice())
@@ -283,7 +306,7 @@ impl MetaTool for SearchToolsTool {
 
         let rank_started = Instant::now();
         let result = crate::services::tool_discovery::ToolDiscoveryService::search(
-            &index,
+            index,
             effective_query,
             server_id_filter,
             detail_level,
@@ -472,7 +495,7 @@ impl MetaTool for SearchToolsTool {
             post_ms,
             accounted_ms,
             unaccounted_ms = total_ms.saturating_sub(accounted_ms),
-            merged_index = index.len(),
+            merged_index = merged_index_len,
             include_inactive,
             scope_all,
             server_id_set,