refactor: scale agentic budget by flat per-doc caps

packages/shared-python/shared/services/retrieval/agentic/core/budget.py

@@ -102,7 +102,7 @@ def __init__(
         total: int,
         planning_ratio: float,
         bootstrap: int = 2000,
-        per_doc_min_share: int = 1500,
+        per_doc_cap: int = 20000,
     ) -> None:
         total = max(int(total), 1)
         bootstrap = max(0, min(int(bootstrap), total))
@@ -117,17 +117,28 @@ def __init__(
             "planning": BudgetPool("planning", planning_capacity),
             "context": BudgetPool("context", context_capacity),
         }
+        self._total = (
+            self._pools["bootstrap"].capacity
+            + self._pools["planning"].capacity
+            + self._pools["context"].capacity
+        )
         self._doc_caps: dict[str, int] = {}
         self._doc_used: dict[str, int] = {}
         self._doc_reserved: dict[str, int] = {}
-        self._per_doc_min_share = max(int(per_doc_min_share), 0)
+        self._planning_ratio = planning_ratio
+        self._per_doc_cap = max(int(per_doc_cap), 0)
         self.total_chunks = 0
         self.total_docs = 0
         self.explored_chunks = 0
         self.explored_docs = 0
         self.trimmed_paths: list[dict[str, Any]] = []
         self._overdraft_events: list[dict[str, Any]] = []
 
+    @property
+    def total(self) -> int:
+        """Spendable token total across bootstrap/planning/context pools."""
+        return self._total
+
     def remaining(self, pool: BudgetPoolName) -> int:
         return self._pools[pool].remaining
 
@@ -138,25 +149,56 @@ def status(self, pool: BudgetPoolName) -> BudgetStatus:
             used_pct=pool_state.used_pct,
         )
 
-    async def allocate_doc_caps(self, doc_chunks: dict[str, int]) -> None:
-        """Allocate planning soft caps by document chunk counts."""
+    async def allocate_doc_caps(self, doc_ids: list[str]) -> None:
+        """Grow pools and total budget for selected docs with a flat per-doc cap.
+
+        After document selection, planning capacity becomes
+        ``len(unique_docs) * per_doc_cap`` (not split by chunk count). Context
+        grows with the same planning/context ratio so evidence budget scales
+        with the selected set. The ledger ``total`` is raised to match the new
+        pool capacities. Each document may spend up to ``per_doc_cap``.
+        """
         async with self._lock:
             self._doc_caps.clear()
             self._doc_used.clear()
             self._doc_reserved.clear()
-            if not doc_chunks:
+            unique_ids = list(dict.fromkeys(doc_id for doc_id in doc_ids if doc_id))
+            if not unique_ids:
                 return
 
-            planning_capacity = self._pools["planning"].capacity
-            total_weight = sum(max(int(count), 1) for count in doc_chunks.values())
-            for doc_id, count in doc_chunks.items():
-                weight = max(int(count), 1)
-                weighted = int(planning_capacity * weight / total_weight)
-                self._doc_caps[doc_id] = min(
-                    planning_capacity,
-                    max(self._per_doc_min_share, weighted),
+            n_docs = len(unique_ids)
+            per_doc = self._per_doc_cap
+            planning_capacity = n_docs * per_doc
+            if self._planning_ratio <= 0:
+                context_capacity = self._pools["context"].capacity
+            elif self._planning_ratio >= 1:
+                context_capacity = 0
+            else:
+                context_capacity = int(
+                    planning_capacity
+                    * (1.0 - self._planning_ratio)
+                    / self._planning_ratio
                 )
 
+            planning = self._pools["planning"]
+            context = self._pools["context"]
+            planning.capacity = max(
+                planning_capacity,
+                planning.used + planning.reserved,
+            )
+            context.capacity = max(
+                context_capacity,
+                context.used + context.reserved,
+            )
+            self._total = (
+                self._pools["bootstrap"].capacity
+                + planning.capacity
+                + context.capacity
+            )
+
+            for doc_id in unique_ids:
+                self._doc_caps[doc_id] = per_doc
+
     async def try_reserve(
         self,
         pool: BudgetPoolName,
@@ -299,6 +341,7 @@ def snapshot(self) -> dict[str, object]:
         if self._overdraft_events:
             snapshot["overdraft_events"] = list(self._overdraft_events)
         snapshot.update({
+            "total": self._total,
             "total_chunks": self.total_chunks,
             "total_docs": self.total_docs,
             "explored_chunks": min(self.explored_chunks, self.total_chunks)

packages/shared-python/shared/services/retrieval/agentic/core/runtime.py

@@ -23,7 +23,7 @@ def build_config_from_env() -> AgentRunConfig:
         token_budget_total=int(os.environ.get("RETRIEVAL_AGENTIC_TOKEN_BUDGET_TOTAL", "40000")),
         planning_ratio=float(os.environ.get("RETRIEVAL_AGENTIC_PLANNING_RATIO", "0.5")),
         bootstrap_budget=int(os.environ.get("RETRIEVAL_AGENTIC_BOOTSTRAP_BUDGET", "2000")),
-        per_doc_min_share=int(os.environ.get("RETRIEVAL_AGENTIC_PER_DOC_MIN_SHARE", "1500")),
+        per_doc_cap=int(os.environ.get("RETRIEVAL_AGENTIC_PER_DOC_CAP", "20000")),
         inventory_aware=os.environ.get("RETRIEVAL_AGENTIC_INVENTORY_AWARE", "true") == "true",
     )
 

packages/shared-python/shared/services/retrieval/agentic/core/types.py

@@ -21,7 +21,7 @@ class AgentRunConfig:
     token_budget_total: int = 40000
     planning_ratio: float = 0.5
     bootstrap_budget: int = 2000
-    per_doc_min_share: int = 1500
+    per_doc_cap: int = 20000
     inventory_aware: bool = True
 
 

packages/shared-python/shared/services/retrieval/agentic/orchestrator.py

@@ -99,9 +99,9 @@ async def run(
             total=config.token_budget_total,
             planning_ratio=config.planning_ratio,
             bootstrap=config.bootstrap_budget,
-            per_doc_min_share=config.per_doc_min_share,
+            per_doc_cap=config.per_doc_cap,
         )
-        total_chunks, total_docs, chunks_count_by_doc = await _load_budget_inventory(
+        total_chunks, total_docs, _chunks_count_by_doc = await _load_budget_inventory(
             db,
             user_id=user_id,
             namespace=namespace,
@@ -267,10 +267,13 @@ async def run(
 
 
         if state.ledger is not None:
-            await state.ledger.allocate_doc_caps({
-                doc.document_id: chunks_count_by_doc.get(doc.document_id, 1)
-                for doc in state.selected_docs
-            })
+            await state.ledger.allocate_doc_caps(
+                [doc.document_id for doc in state.selected_docs]
+            )
+            logger.info(
+                f'agentic: doc caps allocated for {len(state.selected_docs)} docs; '
+                f'token_budget_total={state.ledger.total}'
+            )
 
         # Phase 2 + 3: navigate once, then render evidence for downstream agents.
         evidence_text = ''

packages/shared-python/shared/services/retrieval/workflow/orchestrator.py

@@ -117,7 +117,7 @@ async def run_request(
             total=config.planner_budget,
             planning_ratio=0.0,
             bootstrap=config.planner_budget,
-            per_doc_min_share=0,
+            per_doc_cap=0,
         )
         db_factory = self._get_db_factory()
         async with db_factory() as inventory_db:

packages/shared-python/shared/services/retrieval/workflow/wallet.py

@@ -39,8 +39,8 @@ class BudgetWallet:
     bootstrap_budget: int = field(
         default_factory=lambda: _env_int('RETRIEVAL_AGENTIC_BOOTSTRAP_BUDGET', 2000)
     )
-    per_doc_min_share: int = field(
-        default_factory=lambda: _env_int('RETRIEVAL_AGENTIC_PER_DOC_MIN_SHARE', 1500)
+    per_doc_cap: int = field(
+        default_factory=lambda: _env_int('RETRIEVAL_AGENTIC_PER_DOC_CAP', 20000)
     )
     _allocations: dict[str, int] = field(default_factory=dict, init=False)
     _reclaimed: dict[str, int] = field(default_factory=dict, init=False)
@@ -89,8 +89,14 @@ async def allocate(self, plan: QueryPlan) -> dict[str, BudgetLedger]:
             return dict(self._ledgers)
 
     async def reclaim(self, step_id: str, ledger: BudgetLedger) -> None:
-        """Record unused capacity after a step completes."""
+        """Record unused capacity after a step completes.
+
+        If the step ledger grew its total after document selection
+        (``allocate_doc_caps``), raise this wallet's hard total and the step
+        allocation to match so accounting stays consistent.
+        """
         async with self._lock:
+            self._sync_capacity_growth(step_id, ledger)
             allocated = self._allocations.get(step_id, 0)
             used = self._ledger_used(ledger)
             self._reclaimed[step_id] = max(allocated - used, 0)
@@ -99,19 +105,40 @@ def total_used(self) -> int:
         return sum(self._ledger_used(ledger) for ledger in self._ledgers.values())
 
     def snapshot(self) -> dict[str, object]:
+        allocations, total = self._live_allocations()
+        used = self.total_used()
         return {
-            "total": self.total,
-            "allocated": sum(self._allocations.values()),
-            "used": self.total_used(),
-            "remaining": max(self.total - self.total_used(), 0),
-            "allocations": dict(self._allocations),
+            "total": total,
+            "allocated": sum(allocations.values()),
+            "used": used,
+            "remaining": max(total - used, 0),
+            "allocations": allocations,
             "reclaimed": dict(self._reclaimed),
             "steps": {
                 step_id: ledger.snapshot()
                 for step_id, ledger in self._ledgers.items()
             },
         }
 
+    def _sync_capacity_growth(self, step_id: str, ledger: BudgetLedger) -> None:
+        capacity = max(int(ledger.total), 0)
+        previous = self._allocations.get(step_id, 0)
+        if capacity > previous:
+            self.total += capacity - previous
+            self._allocations[step_id] = capacity
+
+    def _live_allocations(self) -> tuple[dict[str, int], int]:
+        """Return allocations/total including any post-selection ledger growth."""
+        allocations = dict(self._allocations)
+        total = self.total
+        for step_id, ledger in self._ledgers.items():
+            capacity = max(int(ledger.total), 0)
+            previous = allocations.get(step_id, 0)
+            if capacity > previous:
+                total += capacity - previous
+                allocations[step_id] = capacity
+        return allocations, total
+
     def _requested_for_step(self, step: PlannedStep) -> int:
         del step
         return max(self.per_retrieve_step_default, _RETRIEVE_FLOOR)
@@ -122,7 +149,7 @@ def _new_ledger(self, step: PlannedStep, total: int) -> BudgetLedger:
             total=max(total, 1),
             planning_ratio=self.planning_ratio,
             bootstrap=min(self.bootstrap_budget, max(total, 1)),
-            per_doc_min_share=self.per_doc_min_share,
+            per_doc_cap=self.per_doc_cap,
         )
 
     @staticmethod