feat: Phase 2 — Token Budget Manager and model limits

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")

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"
+            ),
+        )