feat(retrieval): add adaptive budget controller for pipeline token management

- Introduce RetrievalBudgetController to track token consumption across all
  pipeline stages (Plan, Search, Evaluate)
- Add BudgetStatus enum with Healthy, Constrained, and Exhausted states
- Implement budget-aware strategy selection in Plan stage
- Apply dynamic beam width adjustment based on budget status in Search stage
- Enable graceful degradation when budget is exhausted
- Track and report budget utilization throughout the pipeline
- Add comprehensive tests for budget controller functionality

Author: zTgx

rust/src/retrieval/mod.rs

@@ -78,8 +78,9 @@ pub use types::StrategyPreference as Strategy;
 
 // Pipeline exports
 pub use pipeline::{
-    CandidateNode, ExecutionGroup, FailurePolicy, PipelineContext, RetrievalMetrics,
-    RetrievalOrchestrator, RetrievalStage, SearchAlgorithm, SearchConfig, StageOutcome,
+    CandidateNode, ExecutionGroup, FailurePolicy, PipelineContext, RetrievalBudgetController,
+    RetrievalMetrics, RetrievalOrchestrator, RetrievalStage, SearchAlgorithm, SearchConfig,
+    StageOutcome, BudgetStatus,
 };
 
 // Re-export PipelineContext as RetrievalContext for stages (alias for clarity)

rust/src/retrieval/pipeline/budget.rs

@@ -0,0 +1,329 @@
+// Copyright (c) 2026 vectorless developers
+// SPDX-License-Identifier: Apache-2.0
+
+//! Adaptive token budget controller for the retrieval pipeline.
+//!
+//! Unlike the Pilot-level [`BudgetController`](crate::retrieval::pilot::BudgetController)
+//! which only tracks Pilot LLM calls, this controller tracks the **entire pipeline's**
+//! token consumption across all stages and provides dynamic budget allocation decisions.
+//!
+//! # Design
+//!
+//! ```text
+//! ┌──────────────────────────────────────────────────┐
+//! │          RetrievalBudgetController                │
+//! │                                                   │
+//! │  total_budget ────────────────────────┬────────── │
+//! │  consumed (from all stages)           │ remaining │
+//! │                                       │           │
+//! │  Plan stage: initial allocation       │           │
+//! │  Search stage: check before iteration │           │
+//! │  Evaluate stage: report & decide      │           │
+//! │  Graceful degradation when low        │           │
+//! └──────────────────────────────────────────────────┘
+//! ```
+
+use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
+use std::sync::Arc;
+
+/// Status of the budget for stage-level decision making.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum BudgetStatus {
+    /// Plenty of budget remaining, proceed normally.
+    Healthy,
+    /// Budget is getting low, consider cheaper strategies.
+    Constrained,
+    /// Budget is exhausted, stop LLM calls and return best results.
+    Exhausted,
+}
+
+impl BudgetStatus {
+    /// Whether LLM calls should still be made.
+    pub fn allow_llm(self) -> bool {
+        matches!(self, Self::Healthy | Self::Constrained)
+    }
+
+    /// Whether the pipeline should stop iterating and return current results.
+    pub fn should_stop(self) -> bool {
+        self == Self::Exhausted
+    }
+}
+
+/// Adaptive budget controller for the retrieval pipeline.
+///
+/// Tracks token consumption across all stages (Plan, Search, Evaluate)
+/// and provides budget-aware decisions for dynamic strategy adjustment.
+///
+/// # Example
+///
+/// ```rust,ignore
+/// let budget = RetrievalBudgetController::new(4000);
+///
+/// // In Search stage: check before starting an iteration
+/// if budget.status().should_stop() {
+///     return StageOutcome::complete(); // graceful degradation
+/// }
+///
+/// // After LLM call: record consumption
+/// budget.record_tokens(350);
+///
+/// // In Evaluate: decide based on remaining budget
+/// if budget.status() == BudgetStatus::Constrained {
+///     // Use cheaper sufficiency check
+/// }
+/// ```
+pub struct RetrievalBudgetController {
+    /// Total token budget for this retrieval operation.
+    total_budget: usize,
+    /// Tokens consumed so far (atomic for thread safety).
+    consumed: AtomicUsize,
+    /// Whether budget exhaustion has been signaled to the pipeline.
+    exhaustion_signaled: AtomicBool,
+    /// Threshold ratio for "constrained" status (e.g. 0.7 = warn at 70% used).
+    constrain_threshold: f32,
+}
+
+// Manual Clone because AtomicUsize/AtomicBool don't impl Clone.
+impl Clone for RetrievalBudgetController {
+    fn clone(&self) -> Self {
+        Self {
+            total_budget: self.total_budget,
+            consumed: AtomicUsize::new(self.consumed.load(Ordering::Relaxed)),
+            exhaustion_signaled: AtomicBool::new(
+                self.exhaustion_signaled.load(Ordering::Relaxed),
+            ),
+            constrain_threshold: self.constrain_threshold,
+        }
+    }
+}
+
+impl RetrievalBudgetController {
+    /// Create a new budget controller with the given total token budget.
+    pub fn new(total_budget: usize) -> Self {
+        Self {
+            total_budget,
+            consumed: AtomicUsize::new(0),
+            exhaustion_signaled: AtomicBool::new(false),
+            constrain_threshold: 0.7,
+        }
+    }
+
+    /// Create with a custom constrain threshold (0.0 - 1.0).
+    ///
+    /// When consumption exceeds `total_budget * threshold`, status becomes Constrained.
+    pub fn with_constrain_threshold(mut self, threshold: f32) -> Self {
+        self.constrain_threshold = threshold.clamp(0.0, 1.0);
+        self
+    }
+
+    /// Get the current budget status.
+    pub fn status(&self) -> BudgetStatus {
+        if self.exhaustion_signaled.load(Ordering::Relaxed) {
+            return BudgetStatus::Exhausted;
+        }
+
+        let consumed = self.consumed.load(Ordering::Relaxed);
+        if consumed >= self.total_budget {
+            self.exhaustion_signaled.store(true, Ordering::Relaxed);
+            return BudgetStatus::Exhausted;
+        }
+
+        let utilization = consumed as f32 / self.total_budget as f32;
+        if utilization >= self.constrain_threshold {
+            BudgetStatus::Constrained
+        } else {
+            BudgetStatus::Healthy
+        }
+    }
+
+    /// Record tokens consumed by any stage.
+    pub fn record_tokens(&self, tokens: usize) {
+        self.consumed.fetch_add(tokens, Ordering::Relaxed);
+    }
+
+    /// Get total tokens consumed so far.
+    pub fn consumed(&self) -> usize {
+        self.consumed.load(Ordering::Relaxed)
+    }
+
+    /// Get remaining token budget.
+    pub fn remaining(&self) -> usize {
+        self.total_budget.saturating_sub(self.consumed.load(Ordering::Relaxed))
+    }
+
+    /// Get total budget.
+    pub fn total_budget(&self) -> usize {
+        self.total_budget
+    }
+
+    /// Get utilization ratio (0.0 - 1.0).
+    pub fn utilization(&self) -> f32 {
+        if self.total_budget == 0 {
+            0.0
+        } else {
+            (self.consumed.load(Ordering::Relaxed) as f32 / self.total_budget as f32).min(1.0)
+        }
+    }
+
+    /// Signal that budget is exhausted (e.g. external trigger).
+    pub fn signal_exhausted(&self) {
+        self.exhaustion_signaled.store(true, Ordering::Relaxed);
+    }
+
+    /// Whether budget exhaustion has been signaled.
+    pub fn is_exhausted(&self) -> bool {
+        self.exhaustion_signaled.load(Ordering::Relaxed)
+            || self.consumed.load(Ordering::Relaxed) >= self.total_budget
+    }
+
+    /// Reset for a new query.
+    pub fn reset(&self) {
+        self.consumed.store(0, Ordering::Relaxed);
+        self.exhaustion_signaled.store(false, Ordering::Relaxed);
+    }
+
+    /// Suggest a search strategy based on budget status and query complexity.
+    ///
+    /// Returns the recommended beam width for the next search iteration.
+    pub fn suggested_beam_width(&self, current_beam: usize, iteration: usize) -> usize {
+        match self.status() {
+            BudgetStatus::Healthy => {
+                // Full power, maybe even increase beam for complex queries
+                current_beam
+            }
+            BudgetStatus::Constrained => {
+                // Reduce beam to save tokens
+                let reduced = if iteration <= 1 { current_beam } else { (current_beam / 2).max(1) };
+                reduced
+            }
+            BudgetStatus::Exhausted => {
+                // No more search iterations worth doing
+                0
+            }
+        }
+    }
+
+    /// Whether another search iteration is worthwhile given budget and confidence.
+    pub fn should_continue_search(&self, current_confidence: f32, iteration: usize) -> bool {
+        if self.is_exhausted() {
+            return false;
+        }
+        // Don't continue if confidence is already good
+        if current_confidence > 0.8 && iteration >= 1 {
+            return false;
+        }
+        // Don't continue if budget is constrained and we have some results
+        if self.status() == BudgetStatus::Constrained && current_confidence > 0.4 {
+            return false;
+        }
+        true
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_budget_healthy() {
+        let budget = RetrievalBudgetController::new(1000);
+        assert_eq!(budget.status(), BudgetStatus::Healthy);
+        assert!(!budget.is_exhausted());
+        assert_eq!(budget.remaining(), 1000);
+    }
+
+    #[test]
+    fn test_budget_constrained() {
+        let budget = RetrievalBudgetController::new(1000);
+        budget.record_tokens(750); // 75% used, above 70% threshold
+        assert_eq!(budget.status(), BudgetStatus::Constrained);
+        assert!(budget.status().allow_llm());
+    }
+
+    #[test]
+    fn test_budget_exhausted() {
+        let budget = RetrievalBudgetController::new(1000);
+        budget.record_tokens(1000);
+        assert_eq!(budget.status(), BudgetStatus::Exhausted);
+        assert!(budget.status().should_stop());
+        assert!(!budget.status().allow_llm());
+    }
+
+    #[test]
+    fn test_budget_exhausted_over() {
+        let budget = RetrievalBudgetController::new(1000);
+        budget.record_tokens(1500);
+        assert_eq!(budget.status(), BudgetStatus::Exhausted);
+    }
+
+    #[test]
+    fn test_budget_signal_exhausted() {
+        let budget = RetrievalBudgetController::new(1000);
+        budget.signal_exhausted();
+        assert_eq!(budget.status(), BudgetStatus::Exhausted);
+        assert_eq!(budget.consumed(), 0); // No tokens actually consumed
+    }
+
+    #[test]
+    fn test_budget_reset() {
+        let budget = RetrievalBudgetController::new(1000);
+        budget.record_tokens(800);
+        assert_eq!(budget.status(), BudgetStatus::Constrained);
+        budget.reset();
+        assert_eq!(budget.status(), BudgetStatus::Healthy);
+        assert_eq!(budget.consumed(), 0);
+    }
+
+    #[test]
+    fn test_suggested_beam_width() {
+        let budget = RetrievalBudgetController::new(1000);
+        // Healthy: keep current beam
+        assert_eq!(budget.suggested_beam_width(4, 0), 4);
+
+        // Constrained: first iteration keeps beam, later reduces
+        budget.record_tokens(750);
+        assert_eq!(budget.suggested_beam_width(4, 0), 4);
+        assert_eq!(budget.suggested_beam_width(4, 2), 2);
+
+        // Exhausted: zero
+        budget.record_tokens(300);
+        assert_eq!(budget.suggested_beam_width(4, 0), 0);
+    }
+
+    #[test]
+    fn test_should_continue_search() {
+        let budget = RetrievalBudgetController::new(1000);
+
+        // Fresh, low confidence: continue
+        assert!(budget.should_continue_search(0.2, 0));
+
+        // High confidence after 1 iteration: stop
+        assert!(!budget.should_continue_search(0.9, 1));
+
+        // Medium confidence, healthy budget: continue
+        assert!(budget.should_continue_search(0.5, 1));
+
+        // Constrained, decent confidence: stop
+        budget.record_tokens(750);
+        assert!(!budget.should_continue_search(0.5, 2));
+
+        // Constrained, low confidence: continue
+        assert!(budget.should_continue_search(0.2, 2));
+    }
+
+    #[test]
+    fn test_utilization() {
+        let budget = RetrievalBudgetController::new(1000);
+        assert!((budget.utilization() - 0.0).abs() < 0.01);
+
+        budget.record_tokens(500);
+        assert!((budget.utilization() - 0.5).abs() < 0.01);
+    }
+
+    #[test]
+    fn test_custom_constrain_threshold() {
+        let budget = RetrievalBudgetController::new(1000).with_constrain_threshold(0.5);
+        budget.record_tokens(500);
+        assert_eq!(budget.status(), BudgetStatus::Constrained);
+    }
+}

rust/src/retrieval/pipeline/context.rs

@@ -12,6 +12,7 @@ use std::time::Instant;
 
 use crate::document::{DocumentTree, NodeId, RetrievalIndex};
 use crate::retrieval::pilot::Pilot;
+use crate::retrieval::pipeline::budget::RetrievalBudgetController;
 use crate::retrieval::types::{
     NavigationDecision, QueryComplexity, ReasoningChain, ReasoningStep, RetrieveOptions,
     RetrieveResponse, SearchPath, StageName, StrategyPreference, SufficiencyLevel,
@@ -201,6 +202,8 @@ pub struct PipelineContext {
     pub options: RetrieveOptions,
     /// Optional Pilot for navigation guidance.
     pub pilot: Option<Arc<dyn Pilot>>,
+    /// Adaptive token budget controller for the entire pipeline.
+    pub budget_controller: RetrievalBudgetController,
 
     // ============ Analyze Stage Output ============
     /// Detected query complexity.
@@ -260,13 +263,15 @@ impl PipelineContext {
     ) -> Self {
         // Build retrieval index for efficient operations
         let retrieval_index = Some(tree.build_retrieval_index());
+        let budget_controller = RetrievalBudgetController::new(options.max_tokens);
 
         Self {
             query: query.into(),
             tree,
             retrieval_index,
             options,
             pilot: None,
+            budget_controller,
             complexity: None,
             keywords: Vec::new(),
             target_sections: Vec::new(),

rust/src/retrieval/pipeline/mod.rs

@@ -44,11 +44,13 @@
 //! let response = orchestrator.execute(tree, query, options).await?;
 //! ```
 
+mod budget;
 mod context;
 mod orchestrator;
 mod outcome;
 mod stage;
 
+pub use budget::{BudgetStatus, RetrievalBudgetController};
 pub use context::{
     CandidateNode, PipelineContext, RetrievalMetrics, SearchAlgorithm, SearchConfig, StageResult,
 };