diff --git a/src/gitmate/config.py b/src/gitmate/config.py index bbfa73f..24f20da 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,9 @@ 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 + template_overhead: int = DEFAULT_TEMPLATE_OVERHEAD class SecretStore(Protocol): @@ -89,6 +107,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 +144,46 @@ 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 + + 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' + 'template_overhead' ({reserved + overhead})" + ) + + return GitmateConfig( + model=model, + commit_style=style, + budget_cap_usd=cap, + ignore_globs=globs, + max_context_tokens=max_ctx, + reserved_output_tokens=reserved, + template_overhead=overhead, + ) def save_config(cfg: GitmateConfig, path: Path | None = None) -> None: @@ -126,4 +195,10 @@ 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 + 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 new file mode 100644 index 0000000..af1a448 --- /dev/null +++ b/src/gitmate/token_budget.py @@ -0,0 +1,276 @@ +"""Token budget manager and token counting abstractions.""" + +from __future__ import annotations + +from dataclasses import dataclass, replace +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 InvalidBudgetError(TokenBudgetError): + """Raised when the calculated effective budget is non-positive.""" + + +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. + + 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 + 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. + + 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, + 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 + 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 be positive ({reserved_output_tokens})." + ) + if default_template_overhead <= 0: + raise InvalidBudgetError( + f"default_template_overhead must be positive ({default_template_overhead})." + ) + 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 + ) + 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( + f"token budget is not positive ({effective_budget}); " + "context_window must exceed reserved_output_tokens + template_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: 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: 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)} + 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) + 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) + 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. + # 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 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( + strategy=BudgetStrategy.NEEDS_CHUNKING, + 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"diff exceeded budget ({current_tokens} > {effective_budget} tokens)" + f"{trunc_info}; requires chunking" + ), + ) diff --git a/tests/test_config.py b/tests/test_config.py index 6e47760..5978e14 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,75 @@ 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, + 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"]) +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() + + +@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 = 2200\nreserved_output_tokens = 2048\n", + encoding="utf-8", + ) + 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 new file mode 100644 index 0000000..ddbe34a --- /dev/null +++ b/tests/test_token_budget.py @@ -0,0 +1,520 @@ +"""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, + InvalidBudgetError, + 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") + + +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(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(InvalidBudgetError, match="not positive"): + mgr.assess([diff]) + with pytest.raises(InvalidBudgetError, match="not positive"): + mgr.assess([]) + + +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 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 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 be positive"): + mgr.assess([diff], template_overhead=0) + with pytest.raises(InvalidBudgetError, match="must be positive"): + 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) + truncated_big = _make_diff( + path="src/service.py", + patch_text=f"[patch omitted for size: +{big.additions}/-{big.deletions} lines]", + ) + 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_big_truncated + 2048 + 500) + + decision = mgr.assess([small, big]) + + assert decision.strategy == BudgetStrategy.TRUNCATED + 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["tests/__snapshots__/app.snap"].patch_text == "s" * 150 + + +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 + + +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([small, med, big]) + + assert decision.strategy == BudgetStrategy.TRUNCATED + # 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/small.py"].patch_text == "s" * 100 + + +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 + + +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