From 9a8953efe981e4e9202c35a456a1423c6c7e531a Mon Sep 17 00:00:00 2001 From: Amar K P Date: Wed, 9 Sep 2026 16:08:46 +0530 Subject: [PATCH 1/6] feat: implement Phase 2 Token Budget Manager --- src/gitmate/config.py | 66 ++++++- src/gitmate/token_budget.py | 245 +++++++++++++++++++++++ tests/test_config.py | 58 ++++++ tests/test_token_budget.py | 376 ++++++++++++++++++++++++++++++++++++ 4 files changed, 744 insertions(+), 1 deletion(-) create mode 100644 src/gitmate/token_budget.py create mode 100644 tests/test_token_budget.py diff --git a/src/gitmate/config.py b/src/gitmate/config.py index bbfa73f..d6b6f3b 100644 --- a/src/gitmate/config.py +++ b/src/gitmate/config.py @@ -20,11 +20,26 @@ DEFAULT_MODEL = "gemini-flash" DEFAULT_COMMIT_STYLE = "conventional" +#: Context window limits for known models; users override via `max_context_tokens`. +DEFAULT_MODEL_CONTEXT_WINDOWS: dict[str, int] = { + "gemini-flash": 1_048_576, + "gemini-3-flash": 1_048_576, + "gemini-3.5-flash-lite": 1_048_576, + "claude-3-5-sonnet": 200_000, + "claude-3-5-haiku": 200_000, +} +DEFAULT_RESERVED_OUTPUT_TOKENS = 2048 +DEFAULT_TEMPLATE_OVERHEAD = 500 + class ConfigError(Exception): """Raised when the config file exists but cannot be parsed.""" +class UnknownModelError(ConfigError): + """Raised when a model has no configured context window.""" + + @dataclass class GitmateConfig: """User settings; the API key is stored separately via SecretStore.""" @@ -33,6 +48,8 @@ class GitmateConfig: commit_style: str = DEFAULT_COMMIT_STYLE budget_cap_usd: float | None = None ignore_globs: list[str] = field(default_factory=list) + max_context_tokens: int | None = None + reserved_output_tokens: int = DEFAULT_RESERVED_OUTPUT_TOKENS class SecretStore(Protocol): @@ -89,6 +106,18 @@ def config_path(base: Path | None = None) -> Path: return config_dir(base) / "config.toml" +def get_context_window(cfg: GitmateConfig) -> int: + """Resolve the model context window limit, checking overrides first.""" + if cfg.max_context_tokens is not None: + return cfg.max_context_tokens + if cfg.model in DEFAULT_MODEL_CONTEXT_WINDOWS: + return DEFAULT_MODEL_CONTEXT_WINDOWS[cfg.model] + raise UnknownModelError( + f"unknown model {cfg.model!r} with no context window configured. " + "Set 'max_context_tokens' in config.toml or use a known model." + ) + + def load_config(path: Path | None = None) -> GitmateConfig: """Load config.toml; a missing file or missing fields fall back to defaults.""" resolved = path if path is not None else config_path() @@ -114,7 +143,38 @@ def load_config(path: Path | None = None) -> GitmateConfig: globs = raw.get("ignore_globs", []) if not isinstance(globs, list) or not all(isinstance(g, str) for g in globs): raise ConfigError(f"cannot parse {resolved}: 'ignore_globs' must be a string list") - return GitmateConfig(model=model, commit_style=style, budget_cap_usd=cap, ignore_globs=globs) + + max_ctx_raw = raw.get("max_context_tokens") + if max_ctx_raw is None: + max_ctx: int | None = None + elif isinstance(max_ctx_raw, bool) or not isinstance(max_ctx_raw, int) or max_ctx_raw <= 0: + raise ConfigError( + f"cannot parse {resolved}: 'max_context_tokens' must be a positive integer" + ) + else: + max_ctx = max_ctx_raw + + reserved_raw = raw.get("reserved_output_tokens", DEFAULT_RESERVED_OUTPUT_TOKENS) + if isinstance(reserved_raw, bool) or not isinstance(reserved_raw, int) or reserved_raw <= 0: + raise ConfigError( + f"cannot parse {resolved}: 'reserved_output_tokens' must be a positive integer" + ) + reserved = reserved_raw + + if max_ctx is not None and max_ctx <= reserved: + raise ConfigError( + f"cannot parse {resolved}: 'max_context_tokens' ({max_ctx}) must be greater " + f"than 'reserved_output_tokens' ({reserved})" + ) + + return GitmateConfig( + model=model, + commit_style=style, + budget_cap_usd=cap, + ignore_globs=globs, + max_context_tokens=max_ctx, + reserved_output_tokens=reserved, + ) def save_config(cfg: GitmateConfig, path: Path | None = None) -> None: @@ -126,4 +186,8 @@ def save_config(cfg: GitmateConfig, path: Path | None = None) -> None: data["budget_cap_usd"] = cfg.budget_cap_usd if cfg.ignore_globs: data["ignore_globs"] = cfg.ignore_globs + if cfg.max_context_tokens is not None: + data["max_context_tokens"] = cfg.max_context_tokens + if cfg.reserved_output_tokens != DEFAULT_RESERVED_OUTPUT_TOKENS: + data["reserved_output_tokens"] = cfg.reserved_output_tokens resolved.write_text(tomli_w.dumps(data), encoding="utf-8") diff --git a/src/gitmate/token_budget.py b/src/gitmate/token_budget.py new file mode 100644 index 0000000..baacc3f --- /dev/null +++ b/src/gitmate/token_budget.py @@ -0,0 +1,245 @@ +"""Token budget manager and token counting abstractions.""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from typing import TYPE_CHECKING, Protocol, runtime_checkable + +from gitmate.config import ( + DEFAULT_MODEL, + DEFAULT_RESERVED_OUTPUT_TOKENS, + DEFAULT_TEMPLATE_OVERHEAD, +) +from gitmate.diff_extractor import FileDiff + +if TYPE_CHECKING: + from google import genai + + +class TokenBudgetError(Exception): + """Base exception for token budgeting and counting errors.""" + + +class TokenCountError(TokenBudgetError): + """Raised when token counting fails.""" + + +class NoApiKeyError(TokenBudgetError): + """Raised when an API key is required but missing.""" + + +@runtime_checkable +class TokenCounter(Protocol): + """Token counting interface so callers and tests can substitute fakes.""" + + def count_tokens(self, text: str, model: str) -> int: + """Count tokens in text for the specified model.""" + ... + + +class GeminiTokenCounter: + """TokenCounter implementation calling the google-genai SDK.""" + + def __init__( + self, + api_key: str | None = None, + client: genai.Client | None = None, + ) -> None: + self._client = client + self._api_key = api_key + + def _get_client(self) -> genai.Client: + if self._client is not None: + return self._client + if not self._api_key: + raise NoApiKeyError( + "Gemini API key is required for token counting. Run 'gitmate config set-key' first." + ) + from google import genai + + self._client = genai.Client(api_key=self._api_key) + return self._client + + def count_tokens(self, text: str, model: str) -> int: + """Count tokens using google-genai client.models.count_tokens.""" + if not text: + return 0 + client = self._get_client() + try: + resp = client.models.count_tokens(model=model, contents=text) + count = getattr(resp, "total_tokens", None) + if count is None: + raise TokenCountError("count_tokens response did not contain 'total_tokens'.") + return int(count) + except Exception as exc: + if isinstance(exc, TokenBudgetError): + raise + raise TokenCountError(f"failed to count tokens via Gemini SDK: {exc}") from exc + + +class BudgetStrategy(str, Enum): + """Strategy decided by the token budget manager.""" + + FITS = "fits" + TRUNCATED = "truncated" + NEEDS_CHUNKING = "needs_chunking" + + +@dataclass +class BudgetDecision: + """Outcome of assessing a diff list against the token budget.""" + + strategy: BudgetStrategy + included_diffs: list[FileDiff] + omitted_diffs: list[FileDiff] + total_tokens: int + budget_limit: int + reserved_output_tokens: int + template_overhead: int + model: str + summary_note: str + + +def format_diff_for_counting(diffs: list[FileDiff]) -> str: + """Format FileDiff objects into text representation for token counting.""" + parts: list[str] = [] + for diff in diffs: + header = f"File: {diff.path} ({diff.status}, +{diff.additions}/-{diff.deletions})" + if diff.old_path: + header += f" (renamed from {diff.old_path})" + if diff.is_binary: + parts.append(f"{header}\n[binary file, diff omitted]") + elif diff.patch_text: + parts.append(f"{header}\n{diff.patch_text}") + else: + parts.append(f"{header}\n[empty or omitted]") + return "\n\n".join(parts) + + +class TokenBudgetManager: + """Manages token allocation, deciding whether diffs fit, truncate, or chunk.""" + + def __init__( + self, + counter: TokenCounter, + model: str = DEFAULT_MODEL, + context_window: int = 1_048_576, + reserved_output_tokens: int = DEFAULT_RESERVED_OUTPUT_TOKENS, + default_template_overhead: int = DEFAULT_TEMPLATE_OVERHEAD, + ) -> None: + self.counter = counter + self.model = model + self.context_window = context_window + self.reserved_output_tokens = reserved_output_tokens + self.default_template_overhead = default_template_overhead + + def count_diff_tokens(self, diffs: list[FileDiff]) -> int: + """Calculate token count for a list of FileDiffs.""" + if not diffs: + return 0 + text = format_diff_for_counting(diffs) + return self.counter.count_tokens(text, self.model) + + def assess( + self, + diffs: list[FileDiff], + template_overhead: int | None = None, + ) -> BudgetDecision: + """Assess diffs against budget: send as-is, truncate, or request chunking.""" + overhead = ( + template_overhead if template_overhead is not None else self.default_template_overhead + ) + effective_budget = self.context_window - self.reserved_output_tokens - overhead + + if not diffs: + return BudgetDecision( + strategy=BudgetStrategy.FITS, + included_diffs=[], + omitted_diffs=[], + total_tokens=0, + budget_limit=effective_budget, + reserved_output_tokens=self.reserved_output_tokens, + template_overhead=overhead, + model=self.model, + summary_note="included 0/0 files (empty diff)", + ) + + # 1. Check if the full diff fits + total_tokens = self.count_diff_tokens(diffs) + if total_tokens <= effective_budget: + return BudgetDecision( + strategy=BudgetStrategy.FITS, + included_diffs=list(diffs), + omitted_diffs=[], + total_tokens=total_tokens, + budget_limit=effective_budget, + reserved_output_tokens=self.reserved_output_tokens, + template_overhead=overhead, + model=self.model, + summary_note=f"included {len(diffs)}/{len(diffs)} files", + ) + + # 2. Over budget: sort candidates by descending patch length to drop largest first + candidates = [d for d in diffs if not d.is_binary and d.patch_text.strip()] + candidates.sort(key=lambda d: len(d.patch_text), reverse=True) + + current_diffs = [ + FileDiff( + path=d.path, + old_path=d.old_path, + status=d.status, + additions=d.additions, + deletions=d.deletions, + patch_text=d.patch_text, + is_binary=d.is_binary, + ) + for d in diffs + ] + diff_index = {d.path: d for d in current_diffs} + omitted: list[FileDiff] = [] + + for candidate in candidates: + note = f"[patch omitted for size: +{candidate.additions}/-{candidate.deletions} lines]" + if len(candidate.patch_text) <= len(note): + continue + omitted.append(candidate) + target = diff_index[candidate.path] + target.patch_text = note + current_tokens = self.count_diff_tokens(current_diffs) + if current_tokens <= effective_budget: + included_count = len(diffs) - len(omitted) + return BudgetDecision( + strategy=BudgetStrategy.TRUNCATED, + included_diffs=current_diffs, + omitted_diffs=omitted, + total_tokens=current_tokens, + budget_limit=effective_budget, + reserved_output_tokens=self.reserved_output_tokens, + template_overhead=overhead, + model=self.model, + summary_note=( + f"included {included_count}/{len(diffs)} files, truncated {len(omitted)}" + ), + ) + + # 3. Even with all candidates truncated, remaining metadata exceeds budget + final_tokens = self.count_diff_tokens(current_diffs) + if omitted: + trunc_info = f" even with all {len(omitted)} files truncated" + else: + trunc_info = "" + return BudgetDecision( + strategy=BudgetStrategy.NEEDS_CHUNKING, + included_diffs=current_diffs, + omitted_diffs=omitted, + total_tokens=final_tokens, + budget_limit=effective_budget, + reserved_output_tokens=self.reserved_output_tokens, + template_overhead=overhead, + model=self.model, + summary_note=( + f"diff exceeded budget ({final_tokens} > {effective_budget} tokens)" + f"{trunc_info}; requires chunking" + ), + ) diff --git a/tests/test_config.py b/tests/test_config.py index 6e47760..3e7460c 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -10,7 +10,9 @@ ConfigError, GitmateConfig, InMemorySecretStore, + UnknownModelError, config_path, + get_context_window, load_config, save_config, ) @@ -70,3 +72,59 @@ def test_ignore_globs_wrong_type_raises(tmp_config_dir: Path) -> None: config_path().write_text('ignore_globs = "*.gen.py"\n', encoding="utf-8") with pytest.raises(ConfigError): load_config() + + +def test_get_context_window_known_models() -> None: + assert get_context_window(GitmateConfig(model="gemini-flash")) == 1_048_576 + assert get_context_window(GitmateConfig(model="claude-3-5-sonnet")) == 200_000 + + +def test_get_context_window_user_override() -> None: + cfg = GitmateConfig(model="gemini-flash", max_context_tokens=50_000) + assert get_context_window(cfg) == 50_000 + + +def test_get_context_window_unknown_model_raises() -> None: + cfg = GitmateConfig(model="unknown-model-xyz") + with pytest.raises(UnknownModelError, match="unknown model 'unknown-model-xyz'"): + get_context_window(cfg) + assert issubclass(UnknownModelError, ConfigError) + + +def test_get_context_window_unknown_model_with_override() -> None: + cfg = GitmateConfig(model="unknown-model-xyz", max_context_tokens=32_000) + assert get_context_window(cfg) == 32_000 + + +def test_budget_config_round_trip(tmp_config_dir: Path) -> None: + cfg = GitmateConfig(max_context_tokens=64_000, reserved_output_tokens=4096) + save_config(cfg) + loaded = load_config() + assert loaded.max_context_tokens == 64_000 + assert loaded.reserved_output_tokens == 4096 + + +@pytest.mark.parametrize("val", ['"big"', "true", "-100", "0"]) +def test_invalid_max_context_tokens_raises(tmp_config_dir: Path, val: str) -> None: + config_path().parent.mkdir(parents=True, exist_ok=True) + config_path().write_text(f"max_context_tokens = {val}\n", encoding="utf-8") + with pytest.raises(ConfigError, match="'max_context_tokens' must be a positive integer"): + load_config() + + +@pytest.mark.parametrize("val", ['"big"', "false", "-50", "0"]) +def test_invalid_reserved_output_tokens_raises(tmp_config_dir: Path, val: str) -> None: + config_path().parent.mkdir(parents=True, exist_ok=True) + config_path().write_text(f"reserved_output_tokens = {val}\n", encoding="utf-8") + with pytest.raises(ConfigError, match="'reserved_output_tokens' must be a positive integer"): + load_config() + + +def test_max_context_tokens_must_exceed_reserved(tmp_config_dir: Path) -> None: + config_path().parent.mkdir(parents=True, exist_ok=True) + config_path().write_text( + "max_context_tokens = 1000\nreserved_output_tokens = 2000\n", + encoding="utf-8", + ) + with pytest.raises(ConfigError, match="must be greater than 'reserved_output_tokens'"): + load_config() diff --git a/tests/test_token_budget.py b/tests/test_token_budget.py new file mode 100644 index 0000000..8d6df79 --- /dev/null +++ b/tests/test_token_budget.py @@ -0,0 +1,376 @@ +"""Tests for TokenBudgetManager, BudgetDecision, and TokenCounter implementations.""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +import pytest + +from gitmate.diff_extractor import FileDiff +from gitmate.token_budget import ( + BudgetDecision, + BudgetStrategy, + GeminiTokenCounter, + NoApiKeyError, + TokenBudgetManager, + TokenCounter, + TokenCountError, + format_diff_for_counting, +) + + +class FakeTokenCounter: + """In-memory TokenCounter for tests; strictly deterministic and offline.""" + + def __init__( + self, + char_rate: int = 1, + fixed_counts: dict[str, int] | None = None, + ) -> None: + self.char_rate = char_rate + self.fixed_counts = fixed_counts or {} + self.calls: list[tuple[str, str]] = [] + + def count_tokens(self, text: str, model: str) -> int: + self.calls.append((text, model)) + if not text: + return 0 + if text in self.fixed_counts: + return self.fixed_counts[text] + return len(text) * self.char_rate + + +def _make_diff( + path: str = "src/main.py", + old_path: str | None = None, + status: str = "modified", + additions: int = 10, + deletions: int = 5, + patch_text: str = "@@ -1,5 +1,10 @@\n+line", + is_binary: bool = False, +) -> FileDiff: + return FileDiff( + path=path, + old_path=old_path, + status=status, # type: ignore[arg-type] + additions=additions, + deletions=deletions, + patch_text=patch_text, + is_binary=is_binary, + ) + + +# --- format_diff_for_counting --- + + +def test_format_diff_for_counting_modified() -> None: + diff = _make_diff( + path="file.py", status="modified", additions=2, deletions=1, patch_text="+a\n-b" + ) + formatted = format_diff_for_counting([diff]) + assert "File: file.py (modified, +2/-1)" in formatted + assert "+a\n-b" in formatted + + +def test_format_diff_for_counting_renamed() -> None: + diff = _make_diff( + path="new.py", + old_path="old.py", + status="renamed", + additions=0, + deletions=0, + patch_text="", + ) + formatted = format_diff_for_counting([diff]) + assert "File: new.py (renamed, +0/-0) (renamed from old.py)" in formatted + assert "[empty or omitted]" in formatted + + +def test_format_diff_for_counting_binary() -> None: + diff = _make_diff(path="image.png", status="added", additions=0, deletions=0, is_binary=True) + formatted = format_diff_for_counting([diff]) + assert "File: image.png (added, +0/-0)" in formatted + assert "[binary file, diff omitted]" in formatted + + +# --- TokenBudgetManager --- + + +def test_empty_diff_fits() -> None: + counter = FakeTokenCounter() + assert isinstance(counter, TokenCounter) + mgr = TokenBudgetManager(counter=counter, context_window=10_000) + decision = mgr.assess([]) + + assert isinstance(decision, BudgetDecision) + assert decision.strategy == BudgetStrategy.FITS + assert decision.included_diffs == [] + assert decision.omitted_diffs == [] + assert decision.total_tokens == 0 + assert decision.budget_limit == 10_000 - 2048 - 500 + assert decision.summary_note == "included 0/0 files (empty diff)" + + +def test_diff_fits_under_budget() -> None: + counter = FakeTokenCounter(char_rate=1) + diff = _make_diff(path="a.py", patch_text="short patch") + mgr = TokenBudgetManager(counter=counter, context_window=10_000) + decision = mgr.assess([diff]) + + assert decision.strategy == BudgetStrategy.FITS + assert len(decision.included_diffs) == 1 + assert decision.included_diffs[0].patch_text == "short patch" + assert decision.omitted_diffs == [] + assert decision.summary_note == "included 1/1 files" + + +def test_exact_boundary_fits() -> None: + diff = _make_diff(path="a.py", patch_text="exact") + formatted = format_diff_for_counting([diff]) + target_tokens = len(formatted) + + # Set context_window so effective_budget == target_tokens + # effective_budget = context_window - reserved(2048) - overhead(500) + context_window = target_tokens + 2048 + 500 + counter = FakeTokenCounter(char_rate=1) + mgr = TokenBudgetManager(counter=counter, context_window=context_window) + + decision = mgr.assess([diff]) + assert decision.strategy == BudgetStrategy.FITS + assert decision.total_tokens == target_tokens + assert decision.budget_limit == target_tokens + + +def test_boundary_plus_one_truncates() -> None: + diff = _make_diff(path="a.py", patch_text="x" * 200) + formatted = format_diff_for_counting([diff]) + # effective_budget is 1 token less than full formatted text + effective_budget = len(formatted) - 1 + context_window = effective_budget + 2048 + 500 + + counter = FakeTokenCounter(char_rate=1) + mgr = TokenBudgetManager(counter=counter, context_window=context_window) + + decision = mgr.assess([diff]) + assert decision.strategy == BudgetStrategy.TRUNCATED + assert len(decision.omitted_diffs) == 1 + assert decision.omitted_diffs[0].path == "a.py" + assert decision.omitted_diffs[0].patch_text == "x" * 200 + assert "[patch omitted for size:" in decision.included_diffs[0].patch_text + assert decision.summary_note == "included 0/1 files, truncated 1" + + +def test_truncation_drops_largest_patch_first() -> None: + # 3 files with different patch sizes: small (50), large (300), medium (150) + small_diff = _make_diff(path="small.py", patch_text="s" * 50) + large_diff = _make_diff(path="large.py", patch_text="L" * 300) + medium_diff = _make_diff(path="med.py", patch_text="m" * 150) + + diffs = [small_diff, large_diff, medium_diff] + + # Calculate token size if large is truncated: + truncated_large_diff = _make_diff( + path="large.py", + patch_text=f"[patch omitted for size: +{large_diff.additions}/-{large_diff.deletions} lines]", + ) + tokens_with_large_truncated = len( + format_diff_for_counting([small_diff, truncated_large_diff, medium_diff]) + ) + + # Set effective budget exactly to accommodate truncating only the largest + context_window = tokens_with_large_truncated + 2048 + 500 + counter = FakeTokenCounter(char_rate=1) + mgr = TokenBudgetManager(counter=counter, context_window=context_window) + + decision = mgr.assess(diffs) + + assert decision.strategy == BudgetStrategy.TRUNCATED + assert len(decision.omitted_diffs) == 1 + assert decision.omitted_diffs[0].path == "large.py" + # Verify original patch preserved in omitted_diffs + assert decision.omitted_diffs[0].patch_text == "L" * 300 + + # Verify small and medium patches remain intact + included_by_path = {d.path: d for d in decision.included_diffs} + assert included_by_path["small.py"].patch_text == "s" * 50 + assert included_by_path["med.py"].patch_text == "m" * 150 + assert "[patch omitted for size:" in included_by_path["large.py"].patch_text + assert decision.summary_note == "included 2/3 files, truncated 1" + + +def test_truncated_diff_preserves_metadata() -> None: + diff = _make_diff( + path="src/service.py", + old_path="src/old_service.py", + status="renamed", + additions=42, + deletions=17, + patch_text="x" * 500, + is_binary=False, + ) + counter = FakeTokenCounter(char_rate=1) + # Budget fits only the truncated note + mgr = TokenBudgetManager(counter=counter, context_window=2048 + 500 + 150) + decision = mgr.assess([diff]) + + assert decision.strategy == BudgetStrategy.TRUNCATED + inc = decision.included_diffs[0] + assert inc.path == "src/service.py" + assert inc.old_path == "src/old_service.py" + assert inc.status == "renamed" + assert inc.additions == 42 + assert inc.deletions == 17 + assert inc.is_binary is False + assert inc.patch_text == "[patch omitted for size: +42/-17 lines]" + + +def test_binary_files_and_empty_patches_not_truncated() -> None: + bin_diff = _make_diff(path="pic.png", is_binary=True, patch_text="") + empty_diff = _make_diff(path="empty.txt", patch_text="") + code_diff = _make_diff(path="code.py", patch_text="c" * 200) + + diffs = [bin_diff, empty_diff, code_diff] + + # Effective budget fits headers + truncated note for code.py + context_window = 2048 + 500 + 200 + counter = FakeTokenCounter(char_rate=1) + mgr = TokenBudgetManager(counter=counter, context_window=context_window) + + decision = mgr.assess(diffs) + assert decision.strategy == BudgetStrategy.TRUNCATED + # Only code_diff was candidate for truncation + assert len(decision.omitted_diffs) == 1 + assert decision.omitted_diffs[0].path == "code.py" + inc_by_path = {d.path: d for d in decision.included_diffs} + assert inc_by_path["pic.png"].is_binary is True + assert inc_by_path["empty.txt"].patch_text == "" + + +def test_extreme_overflow_triggers_needs_chunking() -> None: + diff1 = _make_diff(path="a.py", patch_text="a" * 100) + diff2 = _make_diff(path="b.py", patch_text="b" * 100) + + # Set context window smaller than headers + replacement notes + # So even after dropping all patches, total_tokens > effective_budget + counter = FakeTokenCounter(char_rate=1) + # effective_budget = 10 tokens + mgr = TokenBudgetManager( + counter=counter, + context_window=2048 + 500 + 10, + reserved_output_tokens=2048, + default_template_overhead=500, + ) + + decision = mgr.assess([diff1, diff2]) + assert decision.strategy == BudgetStrategy.NEEDS_CHUNKING + assert len(decision.omitted_diffs) == 2 + assert "requires chunking" in decision.summary_note + # All patches should be replaced with notes in included_diffs + for d in decision.included_diffs: + assert "[patch omitted for size:" in d.patch_text + + +def test_template_overhead_and_reserved_output_deduction() -> None: + diff = _make_diff(path="a.py", patch_text="hello") + counter = FakeTokenCounter(char_rate=1) + mgr = TokenBudgetManager( + counter=counter, + context_window=10_000, + reserved_output_tokens=3_000, + default_template_overhead=1_000, + ) + + decision = mgr.assess([diff], template_overhead=1_500) + assert decision.reserved_output_tokens == 3_000 + assert decision.template_overhead == 1_500 + # effective_budget = 10_000 - 3_000 - 1_500 = 5_500 + assert decision.budget_limit == 5_500 + + +# --- GeminiTokenCounter --- + + +def test_gemini_token_counter_no_api_key_raises() -> None: + counter = GeminiTokenCounter(api_key=None, client=None) + # Empty text returns 0 without raising + assert counter.count_tokens("", "gemini-flash") == 0 + + with pytest.raises(NoApiKeyError, match="Gemini API key is required"): + counter.count_tokens("some text", "gemini-flash") + + +def test_gemini_token_counter_with_mock_client() -> None: + mock_client = MagicMock() + mock_resp = MagicMock() + mock_resp.total_tokens = 77 + mock_client.models.count_tokens.return_value = mock_resp + + counter = GeminiTokenCounter(client=mock_client) + count = counter.count_tokens("hello world", "gemini-flash") + + assert count == 77 + mock_client.models.count_tokens.assert_called_once_with( + model="gemini-flash", + contents="hello world", + ) + + +def test_gemini_token_counter_missing_total_tokens_attribute() -> None: + mock_client = MagicMock() + mock_resp = MagicMock(spec=[]) # no total_tokens attribute + mock_client.models.count_tokens.return_value = mock_resp + + counter = GeminiTokenCounter(client=mock_client) + with pytest.raises(TokenCountError, match="did not contain 'total_tokens'"): + counter.count_tokens("text", "gemini-flash") + + +def test_gemini_token_counter_client_exception_wrapped() -> None: + mock_client = MagicMock() + mock_client.models.count_tokens.side_effect = RuntimeError("API rate limit") + + counter = GeminiTokenCounter(client=mock_client) + with pytest.raises( + TokenCountError, match="failed to count tokens via Gemini SDK: API rate limit" + ): + counter.count_tokens("text", "gemini-flash") + + +def test_candidate_shorter_than_replacement_note_not_truncated() -> None: + # A tiny patch of 3 characters (+x\n) shouldn't be replaced with a 38-char note + tiny_diff = _make_diff(path="tiny.py", patch_text="+x\n") + counter = FakeTokenCounter(char_rate=1) + # Effective budget 10 (less than tiny_diff total formatting) + mgr = TokenBudgetManager( + counter=counter, + context_window=2048 + 500 + 10, + reserved_output_tokens=2048, + default_template_overhead=500, + ) + + decision = mgr.assess([tiny_diff]) + # Cannot be truncated since patch is smaller than note -> triggers chunking + assert decision.strategy == BudgetStrategy.NEEDS_CHUNKING + assert decision.omitted_diffs == [] + assert decision.included_diffs[0].patch_text == "+x\n" + + +def test_gemini_token_counter_lazy_client_init(monkeypatch: pytest.MonkeyPatch) -> None: + mock_client_cls = MagicMock() + mock_instance = MagicMock() + mock_resp = MagicMock() + mock_resp.total_tokens = 15 + mock_instance.models.count_tokens.return_value = mock_resp + mock_client_cls.return_value = mock_instance + + import google.genai + + monkeypatch.setattr(google.genai, "Client", mock_client_cls) + + counter = GeminiTokenCounter(api_key="secret-api-key") + # Client should not be created yet + assert counter._client is None + + result = counter.count_tokens("hello", "gemini-flash") + assert result == 15 + mock_client_cls.assert_called_once_with(api_key="secret-api-key") From 75b399938fb6a5a68d78b884e52c8d02f50d9290 Mon Sep 17 00:00:00 2001 From: Amar K P Date: Wed, 9 Sep 2026 16:24:43 +0530 Subject: [PATCH 2/6] fix: token budget review findings - Reject non-positive effective budgets instead of emitting nonsense limits - Truncate low-signal files first (snapshots, tests, generated), largest first within groups - Working copies via dataclasses.replace; identity-keyed index handles duplicate paths - Reuse last token count instead of a final paid recount - Document BudgetDecision aliasing contract and note semantics --- src/gitmate/token_budget.py | 79 ++++++++++++++++++++++++++----------- tests/test_token_budget.py | 68 +++++++++++++++++++++++++++++++ 2 files changed, 123 insertions(+), 24 deletions(-) diff --git a/src/gitmate/token_budget.py b/src/gitmate/token_budget.py index baacc3f..a179d9d 100644 --- a/src/gitmate/token_budget.py +++ b/src/gitmate/token_budget.py @@ -2,8 +2,9 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, replace from enum import Enum +from pathlib import PurePosixPath from typing import TYPE_CHECKING, Protocol, runtime_checkable from gitmate.config import ( @@ -88,7 +89,15 @@ class BudgetStrategy(str, Enum): @dataclass class BudgetDecision: - """Outcome of assessing a diff list against the token budget.""" + """Outcome of assessing a diff list against the token budget. + + Aliasing contract: entries alias the caller's FileDiff objects, except + truncated replacements, which are fresh copies. In particular + ``omitted_diffs`` holds the originals (full patch preserved). + ``included_diffs`` always carries every input file; truncated ones appear + as one-line note placeholders, so ``len(included_diffs)`` is the file + count while the note counts full patches versus truncations. + """ strategy: BudgetStrategy included_diffs: list[FileDiff] @@ -117,6 +126,32 @@ def format_diff_for_counting(diffs: list[FileDiff]) -> str: return "\n\n".join(parts) +_LOW_SIGNAL_SUFFIXES = (".snap", ".snapshot", ".min.js", ".bundle.js", ".map") +_TEST_DIR_NAMES = frozenset({"test", "tests"}) + + +def _is_low_signal(path: str) -> bool: + """True for files whose full patch is least informative if dropped. + + Test snapshots, test files, and generated bundles carry the lowest + signal-to-noise for a commit message, so their patches go first when + truncating. Char length stays the cost proxy inside and outside this + group: exact per-file token counts would cost one paid API call each. + """ + posix = PurePosixPath(path) + name = posix.name + if name.endswith(_LOW_SIGNAL_SUFFIXES): + return True + if _TEST_DIR_NAMES.intersection(posix.parts): + return True + if name.startswith("test_") or "_test." in name or ".test." in name: + return True + if ".spec." in name or ".generated." in name: + return True + stem = name.rsplit(".", 1)[0] + return stem.endswith(("_generated", ".generated")) + + class TokenBudgetManager: """Manages token allocation, deciding whether diffs fit, truncate, or chunk.""" @@ -151,6 +186,11 @@ def assess( template_overhead if template_overhead is not None else self.default_template_overhead ) effective_budget = self.context_window - self.reserved_output_tokens - overhead + if effective_budget <= 0: + raise ValueError( + f"token budget is not positive ({effective_budget}); " + "context_window must exceed reserved_output_tokens + template_overhead." + ) if not diffs: return BudgetDecision( @@ -180,32 +220,22 @@ def assess( summary_note=f"included {len(diffs)}/{len(diffs)} files", ) - # 2. Over budget: sort candidates by descending patch length to drop largest first + # 2. Over budget: drop low-signal patches first (snapshots, tests, + # generated code), largest first inside and outside that group. candidates = [d for d in diffs if not d.is_binary and d.patch_text.strip()] - candidates.sort(key=lambda d: len(d.patch_text), reverse=True) - - current_diffs = [ - FileDiff( - path=d.path, - old_path=d.old_path, - status=d.status, - additions=d.additions, - deletions=d.deletions, - patch_text=d.patch_text, - is_binary=d.is_binary, - ) - for d in diffs - ] - diff_index = {d.path: d for d in current_diffs} + candidates.sort(key=lambda d: (not _is_low_signal(d.path), -len(d.patch_text))) + + current_diffs = [replace(d) for d in diffs] + working = {id(orig): copy for orig, copy in zip(diffs, current_diffs)} omitted: list[FileDiff] = [] + current_tokens = total_tokens for candidate in candidates: note = f"[patch omitted for size: +{candidate.additions}/-{candidate.deletions} lines]" if len(candidate.patch_text) <= len(note): continue omitted.append(candidate) - target = diff_index[candidate.path] - target.patch_text = note + working[id(candidate)].patch_text = note current_tokens = self.count_diff_tokens(current_diffs) if current_tokens <= effective_budget: included_count = len(diffs) - len(omitted) @@ -223,8 +253,9 @@ def assess( ), ) - # 3. Even with all candidates truncated, remaining metadata exceeds budget - final_tokens = self.count_diff_tokens(current_diffs) + # 3. Even with all candidates truncated, remaining metadata exceeds budget. + # current_tokens already holds the last recount (or the full count when + # nothing was truncated), so no extra paid count call is needed here. if omitted: trunc_info = f" even with all {len(omitted)} files truncated" else: @@ -233,13 +264,13 @@ def assess( strategy=BudgetStrategy.NEEDS_CHUNKING, included_diffs=current_diffs, omitted_diffs=omitted, - total_tokens=final_tokens, + total_tokens=current_tokens, budget_limit=effective_budget, reserved_output_tokens=self.reserved_output_tokens, template_overhead=overhead, model=self.model, summary_note=( - f"diff exceeded budget ({final_tokens} > {effective_budget} tokens)" + f"diff exceeded budget ({current_tokens} > {effective_budget} tokens)" f"{trunc_info}; requires chunking" ), ) diff --git a/tests/test_token_budget.py b/tests/test_token_budget.py index 8d6df79..dfb8953 100644 --- a/tests/test_token_budget.py +++ b/tests/test_token_budget.py @@ -374,3 +374,71 @@ def test_gemini_token_counter_lazy_client_init(monkeypatch: pytest.MonkeyPatch) result = counter.count_tokens("hello", "gemini-flash") assert result == 15 mock_client_cls.assert_called_once_with(api_key="secret-api-key") + + +def test_non_positive_effective_budget_raises() -> None: + counter = FakeTokenCounter() + diff = _make_diff(path="a.py", patch_text="hello") + # reserved alone exceeds the window + mgr = TokenBudgetManager(counter=counter, context_window=1000, reserved_output_tokens=2048) + with pytest.raises(ValueError, match="not positive"): + mgr.assess([diff]) + # per-call overhead can also zero it out + mgr = TokenBudgetManager(counter=counter, context_window=2548) + with pytest.raises(ValueError, match="not positive"): + mgr.assess([diff]) + with pytest.raises(ValueError, match="not positive"): + mgr.assess([]) + + +def test_low_signal_truncated_before_larger_normal() -> None: + normal = _make_diff(path="src/service.py", patch_text="n" * 300) + snap = _make_diff(path="tests/__snapshots__/app.snap", patch_text="s" * 150) + truncated_snap = _make_diff( + path="tests/__snapshots__/app.snap", + patch_text=f"[patch omitted for size: +{snap.additions}/-{snap.deletions} lines]", + ) + tokens_with_snap_truncated = len(format_diff_for_counting([normal, truncated_snap])) + counter = FakeTokenCounter(char_rate=1) + mgr = TokenBudgetManager( + counter=counter, context_window=tokens_with_snap_truncated + 2048 + 500 + ) + + decision = mgr.assess([normal, snap]) + + assert decision.strategy == BudgetStrategy.TRUNCATED + assert [d.path for d in decision.omitted_diffs] == ["tests/__snapshots__/app.snap"] + by_path = {d.path: d for d in decision.included_diffs} + assert by_path["src/service.py"].patch_text == "n" * 300 + + +def test_duplicate_paths_each_processed() -> None: + first = _make_diff(path="dup.py", patch_text="a" * 100) + second = _make_diff(path="dup.py", patch_text="b" * 100) + counter = FakeTokenCounter(char_rate=1) + mgr = TokenBudgetManager(counter=counter, context_window=2048 + 500 + 10) + + decision = mgr.assess([first, second]) + + assert decision.strategy == BudgetStrategy.NEEDS_CHUNKING + assert len(decision.omitted_diffs) == 2 + assert [d.patch_text for d in decision.included_diffs].count( + "[patch omitted for size: +10/-5 lines]" + ) == 2 + # Originals keep their patches; working copies are independent objects. + assert (first.patch_text, second.patch_text) == ("a" * 100, "b" * 100) + + +def test_chunking_path_makes_no_extra_count_call() -> None: + diffs = [ + _make_diff(path="a.py", patch_text="a" * 100), + _make_diff(path="b.py", patch_text="b" * 100), + ] + counter = FakeTokenCounter(char_rate=1) + mgr = TokenBudgetManager(counter=counter, context_window=2048 + 500 + 10) + + decision = mgr.assess(diffs) + + assert decision.strategy == BudgetStrategy.NEEDS_CHUNKING + # 1 initial count + 1 per truncated file; the final total reuses the last one. + assert len(counter.calls) == 3 From 5a9432f4b419a1130cc49c8b3fd624f4dea32412 Mon Sep 17 00:00:00 2001 From: Amar K P Date: Wed, 9 Sep 2026 16:29:32 +0530 Subject: [PATCH 3/6] test: prove low-signal exhaustion order and guard-skipped chunking - Document why ignore_globs and _is_low_signal stay separate mechanisms - Test whole low-signal group truncates before any normal file, in order - Test many guard-skipped small files collectively fall through to chunking --- src/gitmate/token_budget.py | 11 ++++++++++- tests/test_token_budget.py | 36 ++++++++++++++++++++++++++++++++++++ 2 files changed, 46 insertions(+), 1 deletion(-) diff --git a/src/gitmate/token_budget.py b/src/gitmate/token_budget.py index a179d9d..12b24a9 100644 --- a/src/gitmate/token_budget.py +++ b/src/gitmate/token_budget.py @@ -137,6 +137,13 @@ def _is_low_signal(path: str) -> bool: signal-to-noise for a commit message, so their patches go first when truncating. Char length stays the cost proxy inside and outside this group: exact per-file token counts would cost one paid API call each. + + Deliberately separate from ignore_globs: that filter drops noise files + entirely and is user-configured (Phase 1); this only deprioritizes + already-extracted files when over budget, and they are always sent whole + when budget allows. The two lists answer different questions + ("never send" vs "drop first under pressure"), so they stay independent + even where patterns overlap. """ posix = PurePosixPath(path) name = posix.name @@ -221,7 +228,9 @@ def assess( ) # 2. Over budget: drop low-signal patches first (snapshots, tests, - # generated code), largest first inside and outside that group. + # generated code), largest first inside and outside that group. The + # tuple key exhausts the whole low-signal group before any normal + # file is touched, no matter how large the normal files are. candidates = [d for d in diffs if not d.is_binary and d.patch_text.strip()] candidates.sort(key=lambda d: (not _is_low_signal(d.path), -len(d.patch_text))) diff --git a/tests/test_token_budget.py b/tests/test_token_budget.py index dfb8953..3cb74db 100644 --- a/tests/test_token_budget.py +++ b/tests/test_token_budget.py @@ -442,3 +442,39 @@ def test_chunking_path_makes_no_extra_count_call() -> None: assert decision.strategy == BudgetStrategy.NEEDS_CHUNKING # 1 initial count + 1 per truncated file; the final total reuses the last one. assert len(counter.calls) == 3 + + +def test_low_signal_group_exhausted_before_normal_files() -> None: + # Budget forces two truncations: the snap must go first even though the + # normal file is larger, then spilling over to the largest normal file. + big = _make_diff(path="src/big.py", patch_text="n" * 300) + med = _make_diff(path="src/med.py", patch_text="m" * 200) + snap = _make_diff(path="tests/snap.snap", patch_text="s" * 150) + note = f"[patch omitted for size: +{snap.additions}/-{snap.deletions} lines]" + big_trunc = _make_diff(path="src/big.py", patch_text=note) + snap_trunc = _make_diff(path="tests/snap.snap", patch_text=note) + tokens_after_two = len(format_diff_for_counting([big_trunc, med, snap_trunc])) + counter = FakeTokenCounter(char_rate=1) + mgr = TokenBudgetManager(counter=counter, context_window=tokens_after_two + 2048 + 500) + + decision = mgr.assess([big, med, snap]) + + assert decision.strategy == BudgetStrategy.TRUNCATED + # Entire low-signal group first (largest-first within it), then normals. + assert [d.path for d in decision.omitted_diffs] == ["tests/snap.snap", "src/big.py"] + by_path = {d.path: d for d in decision.included_diffs} + assert by_path["src/med.py"].patch_text == "m" * 200 + + +def test_many_guard_skipped_small_files_need_chunking() -> None: + # Each patch is shorter than its replacement note, so the guard skips + # every file — but the headers alone still exceed the budget. + diffs = [_make_diff(path=f"f{i}.py", patch_text="+x\n") for i in range(3)] + counter = FakeTokenCounter(char_rate=1) + mgr = TokenBudgetManager(counter=counter, context_window=2048 + 500 + 10) + + decision = mgr.assess(diffs) + + assert decision.strategy == BudgetStrategy.NEEDS_CHUNKING + assert decision.omitted_diffs == [] + assert [d.patch_text for d in decision.included_diffs] == ["+x\n"] * 3 From 3813496a665f67bb84063a849d42ee0f84c2cc65 Mon Sep 17 00:00:00 2001 From: Amar K P Date: Wed, 9 Sep 2026 16:42:36 +0530 Subject: [PATCH 4/6] fix: adopt size-only truncation policy, wire template_overhead config, and fix mixed chunking note --- src/gitmate/config.py | 15 ++++++- src/gitmate/token_budget.py | 63 ++++++++++-------------------- tests/test_config.py | 24 ++++++++++-- tests/test_token_budget.py | 78 +++++++++++++++++++++++-------------- 4 files changed, 101 insertions(+), 79 deletions(-) diff --git a/src/gitmate/config.py b/src/gitmate/config.py index d6b6f3b..24f20da 100644 --- a/src/gitmate/config.py +++ b/src/gitmate/config.py @@ -50,6 +50,7 @@ class GitmateConfig: ignore_globs: list[str] = field(default_factory=list) max_context_tokens: int | None = None reserved_output_tokens: int = DEFAULT_RESERVED_OUTPUT_TOKENS + template_overhead: int = DEFAULT_TEMPLATE_OVERHEAD class SecretStore(Protocol): @@ -161,10 +162,17 @@ def load_config(path: Path | None = None) -> GitmateConfig: ) reserved = reserved_raw - if max_ctx is not None and max_ctx <= reserved: + overhead_raw = raw.get("template_overhead", DEFAULT_TEMPLATE_OVERHEAD) + if isinstance(overhead_raw, bool) or not isinstance(overhead_raw, int) or overhead_raw <= 0: + raise ConfigError( + f"cannot parse {resolved}: 'template_overhead' must be a positive integer" + ) + overhead = overhead_raw + + if max_ctx is not None and max_ctx <= (reserved + overhead): raise ConfigError( f"cannot parse {resolved}: 'max_context_tokens' ({max_ctx}) must be greater " - f"than 'reserved_output_tokens' ({reserved})" + f"than 'reserved_output_tokens' + 'template_overhead' ({reserved + overhead})" ) return GitmateConfig( @@ -174,6 +182,7 @@ def load_config(path: Path | None = None) -> GitmateConfig: ignore_globs=globs, max_context_tokens=max_ctx, reserved_output_tokens=reserved, + template_overhead=overhead, ) @@ -190,4 +199,6 @@ def save_config(cfg: GitmateConfig, path: Path | None = None) -> None: data["max_context_tokens"] = cfg.max_context_tokens if cfg.reserved_output_tokens != DEFAULT_RESERVED_OUTPUT_TOKENS: data["reserved_output_tokens"] = cfg.reserved_output_tokens + if cfg.template_overhead != DEFAULT_TEMPLATE_OVERHEAD: + data["template_overhead"] = cfg.template_overhead resolved.write_text(tomli_w.dumps(data), encoding="utf-8") diff --git a/src/gitmate/token_budget.py b/src/gitmate/token_budget.py index 12b24a9..c52e811 100644 --- a/src/gitmate/token_budget.py +++ b/src/gitmate/token_budget.py @@ -4,7 +4,6 @@ from dataclasses import dataclass, replace from enum import Enum -from pathlib import PurePosixPath from typing import TYPE_CHECKING, Protocol, runtime_checkable from gitmate.config import ( @@ -22,6 +21,10 @@ class TokenBudgetError(Exception): """Base exception for token budgeting and counting errors.""" +class InvalidBudgetError(TokenBudgetError): + """Raised when the calculated effective budget is non-positive.""" + + class TokenCountError(TokenBudgetError): """Raised when token counting fails.""" @@ -126,41 +129,16 @@ def format_diff_for_counting(diffs: list[FileDiff]) -> str: return "\n\n".join(parts) -_LOW_SIGNAL_SUFFIXES = (".snap", ".snapshot", ".min.js", ".bundle.js", ".map") -_TEST_DIR_NAMES = frozenset({"test", "tests"}) - - -def _is_low_signal(path: str) -> bool: - """True for files whose full patch is least informative if dropped. - - Test snapshots, test files, and generated bundles carry the lowest - signal-to-noise for a commit message, so their patches go first when - truncating. Char length stays the cost proxy inside and outside this - group: exact per-file token counts would cost one paid API call each. - - Deliberately separate from ignore_globs: that filter drops noise files - entirely and is user-configured (Phase 1); this only deprioritizes - already-extracted files when over budget, and they are always sent whole - when budget allows. The two lists answer different questions - ("never send" vs "drop first under pressure"), so they stay independent - even where patterns overlap. - """ - posix = PurePosixPath(path) - name = posix.name - if name.endswith(_LOW_SIGNAL_SUFFIXES): - return True - if _TEST_DIR_NAMES.intersection(posix.parts): - return True - if name.startswith("test_") or "_test." in name or ".test." in name: - return True - if ".spec." in name or ".generated." in name: - return True - stem = name.rsplit(".", 1)[0] - return stem.endswith(("_generated", ".generated")) - - class TokenBudgetManager: - """Manages token allocation, deciding whether diffs fit, truncate, or chunk.""" + """Manages token allocation, deciding whether diffs fit, truncate, or chunk. + + Truncation orders candidate diffs strictly by descending patch length + (largest raw patch first). Signal weighting was evaluated, but common + generated artifacts and lockfiles are already filtered upstream by + ignore_globs; its only non-redundant contribution (test-directory + deprioritization) was actively harmful, stripping critical test diffs in + test-focused commits. Reverted to size-only ordering as the deliberate policy. + """ def __init__( self, @@ -194,7 +172,7 @@ def assess( ) effective_budget = self.context_window - self.reserved_output_tokens - overhead if effective_budget <= 0: - raise ValueError( + raise InvalidBudgetError( f"token budget is not positive ({effective_budget}); " "context_window must exceed reserved_output_tokens + template_overhead." ) @@ -227,12 +205,9 @@ def assess( summary_note=f"included {len(diffs)}/{len(diffs)} files", ) - # 2. Over budget: drop low-signal patches first (snapshots, tests, - # generated code), largest first inside and outside that group. The - # tuple key exhausts the whole low-signal group before any normal - # file is touched, no matter how large the normal files are. + # 2. Over budget: drop largest patches first. candidates = [d for d in diffs if not d.is_binary and d.patch_text.strip()] - candidates.sort(key=lambda d: (not _is_low_signal(d.path), -len(d.patch_text))) + candidates.sort(key=lambda d: len(d.patch_text), reverse=True) current_diffs = [replace(d) for d in diffs] working = {id(orig): copy for orig, copy in zip(diffs, current_diffs)} @@ -265,8 +240,10 @@ def assess( # 3. Even with all candidates truncated, remaining metadata exceeds budget. # current_tokens already holds the last recount (or the full count when # nothing was truncated), so no extra paid count call is needed here. - if omitted: - trunc_info = f" even with all {len(omitted)} files truncated" + if len(omitted) == len(diffs): + trunc_info = f" even with all {len(diffs)} files truncated" + elif omitted: + trunc_info = f" even with {len(omitted)}/{len(diffs)} files truncated" else: trunc_info = "" return BudgetDecision( diff --git a/tests/test_config.py b/tests/test_config.py index 3e7460c..5978e14 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -97,11 +97,16 @@ def test_get_context_window_unknown_model_with_override() -> None: def test_budget_config_round_trip(tmp_config_dir: Path) -> None: - cfg = GitmateConfig(max_context_tokens=64_000, reserved_output_tokens=4096) + cfg = GitmateConfig( + max_context_tokens=64_000, + reserved_output_tokens=4096, + template_overhead=800, + ) save_config(cfg) loaded = load_config() assert loaded.max_context_tokens == 64_000 assert loaded.reserved_output_tokens == 4096 + assert loaded.template_overhead == 800 @pytest.mark.parametrize("val", ['"big"', "true", "-100", "0"]) @@ -120,11 +125,22 @@ def test_invalid_reserved_output_tokens_raises(tmp_config_dir: Path, val: str) - load_config() -def test_max_context_tokens_must_exceed_reserved(tmp_config_dir: Path) -> None: +@pytest.mark.parametrize("val", ['"big"', "false", "-50", "0"]) +def test_invalid_template_overhead_raises(tmp_config_dir: Path, val: str) -> None: + config_path().parent.mkdir(parents=True, exist_ok=True) + config_path().write_text(f"template_overhead = {val}\n", encoding="utf-8") + with pytest.raises(ConfigError, match="'template_overhead' must be a positive integer"): + load_config() + + +def test_max_context_tokens_must_exceed_reserved_plus_overhead(tmp_config_dir: Path) -> None: config_path().parent.mkdir(parents=True, exist_ok=True) config_path().write_text( - "max_context_tokens = 1000\nreserved_output_tokens = 2000\n", + "max_context_tokens = 2200\nreserved_output_tokens = 2048\n", encoding="utf-8", ) - with pytest.raises(ConfigError, match="must be greater than 'reserved_output_tokens'"): + with pytest.raises( + ConfigError, + match="must be greater than 'reserved_output_tokens' \\+ 'template_overhead'", + ): load_config() diff --git a/tests/test_token_budget.py b/tests/test_token_budget.py index 3cb74db..2de97b5 100644 --- a/tests/test_token_budget.py +++ b/tests/test_token_budget.py @@ -11,6 +11,7 @@ BudgetDecision, BudgetStrategy, GeminiTokenCounter, + InvalidBudgetError, NoApiKeyError, TokenBudgetManager, TokenCounter, @@ -381,35 +382,33 @@ def test_non_positive_effective_budget_raises() -> None: diff = _make_diff(path="a.py", patch_text="hello") # reserved alone exceeds the window mgr = TokenBudgetManager(counter=counter, context_window=1000, reserved_output_tokens=2048) - with pytest.raises(ValueError, match="not positive"): + with pytest.raises(InvalidBudgetError, match="not positive"): mgr.assess([diff]) # per-call overhead can also zero it out mgr = TokenBudgetManager(counter=counter, context_window=2548) - with pytest.raises(ValueError, match="not positive"): + with pytest.raises(InvalidBudgetError, match="not positive"): mgr.assess([diff]) - with pytest.raises(ValueError, match="not positive"): + with pytest.raises(InvalidBudgetError, match="not positive"): mgr.assess([]) -def test_low_signal_truncated_before_larger_normal() -> None: - normal = _make_diff(path="src/service.py", patch_text="n" * 300) - snap = _make_diff(path="tests/__snapshots__/app.snap", patch_text="s" * 150) - truncated_snap = _make_diff( - path="tests/__snapshots__/app.snap", - patch_text=f"[patch omitted for size: +{snap.additions}/-{snap.deletions} lines]", +def test_largest_patch_truncated_first() -> None: + big = _make_diff(path="src/service.py", patch_text="n" * 300) + small = _make_diff(path="tests/__snapshots__/app.snap", patch_text="s" * 150) + truncated_big = _make_diff( + path="src/service.py", + patch_text=f"[patch omitted for size: +{big.additions}/-{big.deletions} lines]", ) - tokens_with_snap_truncated = len(format_diff_for_counting([normal, truncated_snap])) + tokens_with_big_truncated = len(format_diff_for_counting([truncated_big, small])) counter = FakeTokenCounter(char_rate=1) - mgr = TokenBudgetManager( - counter=counter, context_window=tokens_with_snap_truncated + 2048 + 500 - ) + mgr = TokenBudgetManager(counter=counter, context_window=tokens_with_big_truncated + 2048 + 500) - decision = mgr.assess([normal, snap]) + decision = mgr.assess([small, big]) assert decision.strategy == BudgetStrategy.TRUNCATED - assert [d.path for d in decision.omitted_diffs] == ["tests/__snapshots__/app.snap"] + assert [d.path for d in decision.omitted_diffs] == ["src/service.py"] by_path = {d.path: d for d in decision.included_diffs} - assert by_path["src/service.py"].patch_text == "n" * 300 + assert by_path["tests/__snapshots__/app.snap"].patch_text == "s" * 150 def test_duplicate_paths_each_processed() -> None: @@ -444,26 +443,27 @@ def test_chunking_path_makes_no_extra_count_call() -> None: assert len(counter.calls) == 3 -def test_low_signal_group_exhausted_before_normal_files() -> None: - # Budget forces two truncations: the snap must go first even though the - # normal file is larger, then spilling over to the largest normal file. - big = _make_diff(path="src/big.py", patch_text="n" * 300) - med = _make_diff(path="src/med.py", patch_text="m" * 200) - snap = _make_diff(path="tests/snap.snap", patch_text="s" * 150) - note = f"[patch omitted for size: +{snap.additions}/-{snap.deletions} lines]" - big_trunc = _make_diff(path="src/big.py", patch_text=note) - snap_trunc = _make_diff(path="tests/snap.snap", patch_text=note) - tokens_after_two = len(format_diff_for_counting([big_trunc, med, snap_trunc])) +def test_candidates_truncated_strictly_in_descending_size_order() -> None: + # Three candidate diffs of descending patch size: 400, 250, 100 chars. + # When budget allows only 1 file intact, the 400 and 250 files are truncated in order. + big = _make_diff(path="src/big.py", patch_text="b" * 400) + med = _make_diff(path="src/med.py", patch_text="m" * 250) + small = _make_diff(path="src/small.py", patch_text="s" * 100) + note_big = f"[patch omitted for size: +{big.additions}/-{big.deletions} lines]" + note_med = f"[patch omitted for size: +{med.additions}/-{med.deletions} lines]" + big_trunc = _make_diff(path="src/big.py", patch_text=note_big) + med_trunc = _make_diff(path="src/med.py", patch_text=note_med) + tokens_after_two = len(format_diff_for_counting([big_trunc, med_trunc, small])) counter = FakeTokenCounter(char_rate=1) mgr = TokenBudgetManager(counter=counter, context_window=tokens_after_two + 2048 + 500) - decision = mgr.assess([big, med, snap]) + decision = mgr.assess([small, med, big]) assert decision.strategy == BudgetStrategy.TRUNCATED - # Entire low-signal group first (largest-first within it), then normals. - assert [d.path for d in decision.omitted_diffs] == ["tests/snap.snap", "src/big.py"] + # Largest patches dropped first: big (400), then med (250). + assert [d.path for d in decision.omitted_diffs] == ["src/big.py", "src/med.py"] by_path = {d.path: d for d in decision.included_diffs} - assert by_path["src/med.py"].patch_text == "m" * 200 + assert by_path["src/small.py"].patch_text == "s" * 100 def test_many_guard_skipped_small_files_need_chunking() -> None: @@ -478,3 +478,21 @@ def test_many_guard_skipped_small_files_need_chunking() -> None: assert decision.strategy == BudgetStrategy.NEEDS_CHUNKING assert decision.omitted_diffs == [] assert [d.patch_text for d in decision.included_diffs] == ["+x\n"] * 3 + + +def test_mixed_truncated_and_guarded_diffs_trigger_needs_chunking() -> None: + # 1 large file is truncated, but 2 small files are guard-skipped and remaining + # diff still exceeds the effective budget. + big = _make_diff(path="big.py", patch_text="b" * 300) + small1 = _make_diff(path="s1.py", patch_text="+x\n") + small2 = _make_diff(path="s2.py", patch_text="+y\n") + counter = FakeTokenCounter(char_rate=1) + # Effective budget 10: even with big truncated, remaining tokens exceed 10 + mgr = TokenBudgetManager(counter=counter, context_window=2048 + 500 + 10) + + decision = mgr.assess([big, small1, small2]) + + assert decision.strategy == BudgetStrategy.NEEDS_CHUNKING + assert [d.path for d in decision.omitted_diffs] == ["big.py"] + # Ratio format: 1 of 3 files truncated, NOT "even with all 1 files truncated" + assert "even with 1/3 files truncated; requires chunking" in decision.summary_note From 1214f61e1844cc4e2b06ebc1512a06f1c02993b5 Mon Sep 17 00:00:00 2001 From: Amar K P Date: Wed, 9 Sep 2026 16:50:24 +0530 Subject: [PATCH 5/6] fix: validate negative budget inputs, correct aliasing docs - Reject negative reserved_output_tokens, default_template_overhead, and per-call template_overhead - BudgetDecision docstring now matches copy-all implementation on truncate paths --- src/gitmate/token_budget.py | 22 ++++++++++++++++------ tests/test_token_budget.py | 12 ++++++++++++ 2 files changed, 28 insertions(+), 6 deletions(-) diff --git a/src/gitmate/token_budget.py b/src/gitmate/token_budget.py index c52e811..59764af 100644 --- a/src/gitmate/token_budget.py +++ b/src/gitmate/token_budget.py @@ -94,12 +94,12 @@ class BudgetStrategy(str, Enum): class BudgetDecision: """Outcome of assessing a diff list against the token budget. - Aliasing contract: entries alias the caller's FileDiff objects, except - truncated replacements, which are fresh copies. In particular - ``omitted_diffs`` holds the originals (full patch preserved). - ``included_diffs`` always carries every input file; truncated ones appear - as one-line note placeholders, so ``len(included_diffs)`` is the file - count while the note counts full patches versus truncations. + Aliasing contract: on the FITS path entries alias the caller's objects. + On truncate/chunk paths ``included_diffs`` are fresh copies (truncated + ones carry note placeholders); ``omitted_diffs`` always holds the + original objects with full patches preserved. ``included_diffs`` always + carries every input file, so ``len(included_diffs)`` is the file count + while the note counts full patches versus truncations. """ strategy: BudgetStrategy @@ -151,6 +151,14 @@ def __init__( self.counter = counter self.model = model self.context_window = context_window + if reserved_output_tokens < 0: + raise InvalidBudgetError( + f"reserved_output_tokens must not be negative ({reserved_output_tokens})." + ) + if default_template_overhead < 0: + raise InvalidBudgetError( + f"default_template_overhead must not be negative ({default_template_overhead})." + ) self.reserved_output_tokens = reserved_output_tokens self.default_template_overhead = default_template_overhead @@ -170,6 +178,8 @@ def assess( overhead = ( template_overhead if template_overhead is not None else self.default_template_overhead ) + if overhead < 0: + raise InvalidBudgetError(f"template_overhead must not be negative ({overhead}).") effective_budget = self.context_window - self.reserved_output_tokens - overhead if effective_budget <= 0: raise InvalidBudgetError( diff --git a/tests/test_token_budget.py b/tests/test_token_budget.py index 2de97b5..893cc51 100644 --- a/tests/test_token_budget.py +++ b/tests/test_token_budget.py @@ -392,6 +392,18 @@ def test_non_positive_effective_budget_raises() -> None: mgr.assess([]) +def test_negative_budget_inputs_raise() -> None: + counter = FakeTokenCounter() + diff = _make_diff(path="a.py", patch_text="hello") + with pytest.raises(InvalidBudgetError, match="must not be negative"): + TokenBudgetManager(counter=counter, reserved_output_tokens=-5) + with pytest.raises(InvalidBudgetError, match="must not be negative"): + TokenBudgetManager(counter=counter, default_template_overhead=-5) + mgr = TokenBudgetManager(counter=counter, context_window=10_000) + with pytest.raises(InvalidBudgetError, match="must not be negative"): + mgr.assess([diff], template_overhead=-100) + + def test_largest_patch_truncated_first() -> None: big = _make_diff(path="src/service.py", patch_text="n" * 300) small = _make_diff(path="tests/__snapshots__/app.snap", patch_text="s" * 150) From 8edc4739f006e5041f88f5bd087384d9342e5d9c Mon Sep 17 00:00:00 2001 From: Amar K P Date: Wed, 9 Sep 2026 16:54:19 +0530 Subject: [PATCH 6/6] fix: enforce positive budget inputs and clarify BudgetDecision aliasing docstring --- src/gitmate/token_budget.py | 28 ++++++++++++++++------------ tests/test_token_budget.py | 18 ++++++++++++++---- 2 files changed, 30 insertions(+), 16 deletions(-) diff --git a/src/gitmate/token_budget.py b/src/gitmate/token_budget.py index 59764af..af1a448 100644 --- a/src/gitmate/token_budget.py +++ b/src/gitmate/token_budget.py @@ -94,12 +94,14 @@ class BudgetStrategy(str, Enum): class BudgetDecision: """Outcome of assessing a diff list against the token budget. - Aliasing contract: on the FITS path entries alias the caller's objects. - On truncate/chunk paths ``included_diffs`` are fresh copies (truncated - ones carry note placeholders); ``omitted_diffs`` always holds the - original objects with full patches preserved. ``included_diffs`` always - carries every input file, so ``len(included_diffs)`` is the file count - while the note counts full patches versus truncations. + Aliasing contract: on the FITS path entries in ``included_diffs`` alias + the caller's objects unchanged. On truncate and chunk paths, all entries + in ``included_diffs`` are fresh copies (via ``replace()``) so caller + objects are never mutated in-place; truncated files carry one-line note + placeholders. ``omitted_diffs`` always holds the original objects with full + patches preserved. ``included_diffs`` always carries every input file, so + ``len(included_diffs)`` is the file count while the note counts full patches + versus truncations. """ strategy: BudgetStrategy @@ -151,13 +153,15 @@ def __init__( self.counter = counter self.model = model self.context_window = context_window - if reserved_output_tokens < 0: + if context_window <= 0: + raise InvalidBudgetError(f"context_window must be positive ({context_window}).") + if reserved_output_tokens <= 0: raise InvalidBudgetError( - f"reserved_output_tokens must not be negative ({reserved_output_tokens})." + f"reserved_output_tokens must be positive ({reserved_output_tokens})." ) - if default_template_overhead < 0: + if default_template_overhead <= 0: raise InvalidBudgetError( - f"default_template_overhead must not be negative ({default_template_overhead})." + f"default_template_overhead must be positive ({default_template_overhead})." ) self.reserved_output_tokens = reserved_output_tokens self.default_template_overhead = default_template_overhead @@ -178,8 +182,8 @@ def assess( overhead = ( template_overhead if template_overhead is not None else self.default_template_overhead ) - if overhead < 0: - raise InvalidBudgetError(f"template_overhead must not be negative ({overhead}).") + if overhead <= 0: + raise InvalidBudgetError(f"template_overhead must be positive ({overhead}).") effective_budget = self.context_window - self.reserved_output_tokens - overhead if effective_budget <= 0: raise InvalidBudgetError( diff --git a/tests/test_token_budget.py b/tests/test_token_budget.py index 893cc51..ddbe34a 100644 --- a/tests/test_token_budget.py +++ b/tests/test_token_budget.py @@ -392,15 +392,25 @@ def test_non_positive_effective_budget_raises() -> None: mgr.assess([]) -def test_negative_budget_inputs_raise() -> None: +def test_non_positive_budget_inputs_raise() -> None: counter = FakeTokenCounter() diff = _make_diff(path="a.py", patch_text="hello") - with pytest.raises(InvalidBudgetError, match="must not be negative"): + with pytest.raises(InvalidBudgetError, match="must be positive"): + TokenBudgetManager(counter=counter, context_window=0) + with pytest.raises(InvalidBudgetError, match="must be positive"): + TokenBudgetManager(counter=counter, context_window=-100) + with pytest.raises(InvalidBudgetError, match="must be positive"): + TokenBudgetManager(counter=counter, reserved_output_tokens=0) + with pytest.raises(InvalidBudgetError, match="must be positive"): TokenBudgetManager(counter=counter, reserved_output_tokens=-5) - with pytest.raises(InvalidBudgetError, match="must not be negative"): + with pytest.raises(InvalidBudgetError, match="must be positive"): + TokenBudgetManager(counter=counter, default_template_overhead=0) + with pytest.raises(InvalidBudgetError, match="must be positive"): TokenBudgetManager(counter=counter, default_template_overhead=-5) mgr = TokenBudgetManager(counter=counter, context_window=10_000) - with pytest.raises(InvalidBudgetError, match="must not be negative"): + with pytest.raises(InvalidBudgetError, match="must be positive"): + mgr.assess([diff], template_overhead=0) + with pytest.raises(InvalidBudgetError, match="must be positive"): mgr.assess([diff], template_overhead=-100)