From 7864ab9af06fffe215f73024271fcf1672f9cc6f Mon Sep 17 00:00:00 2001
From: Salil Das <11658960+sadlilas@users.noreply.github.com>
Date: Sat, 21 Feb 2026 12:28:31 -0800
Subject: [PATCH] feat: respect provider context_budget_cap for cost-optimal
 compaction
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Read context_budget_cap from provider defaults and cap the computed
budget. Providers with pricing cliffs (Anthropic, Gemini) set this to
keep sessions in the standard pricing zone while preserving the full
context window as a safety net. Providers without pricing cliffs
(OpenAI) don't set the key and are unaffected.

Applied in both the get_model_info() and get_info().defaults code paths.

Replaces the fraction-based approach (PR #5) which applied uniformly to
all providers regardless of their pricing model.

Related: microsoft-amplifier/amplifier-support#57

🤖 Generated with [Amplifier](https://github.com/microsoft/amplifier)

Co-Authored-By: Amplifier <240397093+microsoft-amplifier@users.noreply.github.com>
---
 amplifier_module_context_simple/__init__.py | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/amplifier_module_context_simple/__init__.py b/amplifier_module_context_simple/__init__.py
index 5b856b3..e1d3269 100644
--- a/amplifier_module_context_simple/__init__.py
+++ b/amplifier_module_context_simple/__init__.py
@@ -1183,6 +1183,16 @@ def _calculate_budget(self, token_budget: int | None, provider: Any | None) -> i
                         if context_window and max_output:
                             reserved_output = int(max_output * output_reserve_fraction)
                             budget = context_window - reserved_output - safety_margin
+                            # Apply provider's cost-optimal budget cap if set
+                            budget_cap = (
+                                provider.get_info().defaults or {}
+                            ).get("context_budget_cap")
+                            if budget_cap is not None:
+                                budget = min(budget, budget_cap - safety_margin)
+                                logger.info(
+                                    f"Budget capped by provider hint: {budget:,} "
+                                    f"(cap={budget_cap:,})"
+                                )
                             logger.info(
                                 f"Budget from provider model info: {budget:,} "
                                 f"(context={context_window:,}, reserved_output={reserved_output:,} "
@@ -1199,6 +1209,14 @@ def _calculate_budget(self, token_budget: int | None, provider: Any | None) -> i
                 if context_window and max_output_tokens:
                     reserved_output = int(max_output_tokens * output_reserve_fraction)
                     budget = context_window - reserved_output - safety_margin
+                    # Apply provider's cost-optimal budget cap if set
+                    budget_cap = defaults.get("context_budget_cap")
+                    if budget_cap is not None:
+                        budget = min(budget, budget_cap - safety_margin)
+                        logger.info(
+                            f"Budget capped by provider hint: {budget:,} "
+                            f"(cap={budget_cap:,})"
+                        )
                     logger.info(
                         f"Budget from provider defaults: {budget:,} "
                         f"(context={context_window:,}, reserved_output={reserved_output:,} "