Coordinate foreground search and background inference CPU budgets

crates/searchi-service/src/lib.rs

@@ -1,6 +1,10 @@
 //! Shared offline hybrid search using real Google embeddings and persisted,
 //! filtered HNSW retrieval. Production-scale validation remains a release gate.
+mod coordination;
 mod media;
+pub use coordination::{
+    BudgetedProvider, InferenceBudget, InferencePermit, InferenceRole, PendingForeground,
+};
 mod prepared;
 pub use prepared::PreparedGraphs;
 mod pipeline;
@@ -52,6 +56,11 @@ fn reusable_policy(expected: &str, active: &str) -> bool {
 /// Model-specific query/document prompts and actual tokenizer windows belong
 /// to the provider. Storage and ranking never assume a particular dimension.
 pub trait EmbeddingProvider: Send {
+    /// Maximum document batch admitted as one inference call. Coordinated
+    /// background sessions yield at every document boundary.
+    fn document_batch_limit(&self) -> usize {
+        MAX_BATCH
+    }
     fn fingerprint(&self) -> &str;
     fn dimensions(&self) -> usize;
     fn split_windows(&mut self, text: &str, title: &str) -> Result<Vec<TextWindow>>;
@@ -66,6 +75,32 @@ pub trait EmbeddingProvider: Send {
     fn embed_query_debug(&mut self, query: &str) -> Result<Vec<f32>> {
         self.embed_query(query)
     }
+    fn embed_query_traced(
+        &mut self,
+        query: &str,
+        debug: bool,
+        _trace: &mut TraceCollector,
+        _parent: Option<SpanToken>,
+    ) -> Result<Vec<f32>> {
+        if debug {
+            self.embed_query_debug(query)
+        } else {
+            self.embed_query(query)
+        }
+    }
+    fn embed_code_query_traced(
+        &mut self,
+        query: &str,
+        debug: bool,
+        _trace: &mut TraceCollector,
+        _parent: Option<SpanToken>,
+    ) -> Result<Vec<f32>> {
+        if debug {
+            self.embed_code_query_debug(query)
+        } else {
+            self.embed_code_query(query)
+        }
+    }
     fn query_diagnostics(&self) -> std::collections::BTreeMap<String, u64> {
         Default::default()
     }
@@ -174,6 +209,21 @@ pub struct HybridSearch {
 }
 
 impl HybridSearch {
+    /// Attach shared admission without replacing the warmed model or query LRU.
+    pub fn with_budget(
+        mut self,
+        budget: &InferenceBudget,
+        role: InferenceRole,
+        actual_worker_threads: usize,
+    ) -> Self {
+        self.embedder = Box::new(BudgetedProvider::new(
+            self.embedder,
+            budget,
+            role,
+            actual_worker_threads,
+        ));
+        self
+    }
     /// Connect background publication to an existing warmed reader.
     pub fn share_prepared_graphs(&mut self, graphs: &PreparedGraphs) -> Result<()> {
         self.store.share_prepared_graphs(graphs)
@@ -442,11 +492,12 @@ impl HybridSearch {
         trace.counter(span, "query_cache_hit", 0);
         trace.counter(span, "query_cache_miss", 1);
         let vector = (|| {
-            let vector = match (code, trace.is_debug()) {
-                (true, true) => self.embedder.embed_code_query_debug(query)?,
-                (true, false) => self.embedder.embed_code_query(query)?,
-                (false, true) => self.embedder.embed_query_debug(query)?,
-                (false, false) => self.embedder.embed_query(query)?,
+            let vector = if code {
+                self.embedder
+                    .embed_code_query_traced(query, trace.is_debug(), trace, span)?
+            } else {
+                self.embedder
+                    .embed_query_traced(query, trace.is_debug(), trace, span)?
             };
             ensure!(
                 vector.len() == self.embedder.dimensions()