From 1b238e903d659bb18d62fffaeece7e2ff516f1d8 Mon Sep 17 00:00:00 2001 From: De-pitcher Date: Mon, 7 Sep 2026 05:57:06 +0100 Subject: [PATCH 1/4] feat(phase35): token budget governor, adaptive context compression, and budget commands --- agentcli/cli.py | 54 +++- agentcli/memory/__init__.py | 20 +- agentcli/memory/adaptive_compressor.py | 291 ++++++++++++++++++++ agentcli/memory/governor.py | 365 +++++++++++++++++++++++++ agentcli/session.py | 44 ++- agentcli/ui/prompt.py | 8 +- agentcli/ui/tui_app.py | 45 +++ tests/test_phase35_token_governor.py | 341 +++++++++++++++++++++++ 8 files changed, 1145 insertions(+), 23 deletions(-) create mode 100644 agentcli/memory/adaptive_compressor.py create mode 100644 agentcli/memory/governor.py create mode 100644 tests/test_phase35_token_governor.py diff --git a/agentcli/cli.py b/agentcli/cli.py index 5e053c8..f7ae3a0 100644 --- a/agentcli/cli.py +++ b/agentcli/cli.py @@ -681,20 +681,51 @@ async def run_chat(args: argparse.Namespace, config: Config) -> int: print("-----------------------------------------------\n") continue - if user_input.startswith("/budget"): - parts = user_input.split(maxsplit=1) - if len(parts) == 1: + if user_input.startswith(("/budget", "\\budget")): + budget_parts = user_input.split(maxsplit=2) + budget_subcmd = budget_parts[1].lower() if len(budget_parts) > 1 else "" + budget_val = budget_parts[2].strip() if len(budget_parts) > 2 else "" + + if not budget_subcmd or budget_subcmd in ("status", "info", "show"): current_tier = config.routing.budget_tier print(f"Current budget tier: {current_tier}") - else: - tier = parts[1].strip().lower() - if tier in {"low", "medium", "high"}: - config.routing.budget_tier = tier + print("\n" + session.governor.format_summary() + "\n") + elif budget_subcmd in ("low", "medium", "high", "tier"): + tier_name = budget_val if budget_subcmd == "tier" else budget_subcmd + if tier_name in ("low", "medium", "high"): + config.routing.budget_tier = tier_name if session.router is not None: - session.router._budget_tier = tier - print(f"Budget tier updated to: {tier}") + session.router._budget_tier = tier_name + print(f"Budget tier updated to: {tier_name}") + else: + print(f"Invalid budget tier '{tier_name}'. Choose from: low, medium, high") + elif budget_subcmd in ("set", "limit", "max-cost"): + if not budget_val: + print("Usage: /budget set ") else: - print(f"Invalid budget tier '{tier}'. Choose from: low, medium, high") + try: + val = float(budget_val.lstrip("$")) + session.governor.set_budget(max_cost_usd=val) + config.routing.max_cost_usd = val + print(f"Session budget ceiling set to: ${val:.4f} USD") + except ValueError: + print(f"Invalid budget amount: '{budget_val}'. Must be a positive number.") + elif budget_subcmd in ("tokens", "max-tokens"): + if not budget_val: + print("Usage: /budget max-tokens ") + else: + try: + t_val = int(budget_val.replace(",", "")) + session.governor.set_budget(max_tokens=t_val) + print(f"Session token budget ceiling set to: {t_val:,} tokens") + except ValueError: + print(f"Invalid token count: '{budget_val}'.") + elif budget_subcmd in ("reset", "clear"): + session.governor.reset() + session.cumulative_cost_usd = 0.0 + print("Budget and cost counters reset for current session.") + else: + print(f"Invalid budget tier '{budget_subcmd}'. Choose from: low, medium, high") continue if user_input in {"/models", "/model"} or user_input.startswith( @@ -812,7 +843,8 @@ async def run_chat(args: argparse.Namespace, config: Config) -> int: f"Token Usage: {stats['total_tokens']} total " f"({stats['user_tokens']} prompt, {stats['assistant_tokens']} completion)" ) - print(f"Estimated Cost: ${cost:.6f} USD") + print(f"Estimated Cost: ${cost:.6f} USD\n") + print(session.governor.format_summary()) continue if user_input in {"/diff", "/diffs"}: diff --git a/agentcli/memory/__init__.py b/agentcli/memory/__init__.py index d7c53b6..f4bae69 100644 --- a/agentcli/memory/__init__.py +++ b/agentcli/memory/__init__.py @@ -1,7 +1,13 @@ -"""agentcli memory package — conversation persistence, context caching, and budgeting (Phase 5).""" +"""agentcli memory package — conversation persistence, context caching, budgeting, and adaptive governor (Phase 35).""" from __future__ import annotations +from .adaptive_compressor import ( + AdaptiveContextCompressor, + CompressionMetrics, + collapse_repeated_lines, + strip_ansi_codes, +) from .budget import ( CHARS_PER_TOKEN, DEFAULT_BUDGET_RATIO, @@ -13,6 +19,11 @@ ) from .cache import CachedFileContext, ContextCache, get_default_context_cache from .context_pool import ContextItem, SharedContextPool +from .governor import ( + BudgetHealth, + TokenBudgetGovernor, + UsageRecord, +) from .store import ( MemoryStore, MessageRecord, @@ -24,17 +35,24 @@ "CHARS_PER_TOKEN", "DEFAULT_BUDGET_RATIO", "DEFAULT_CONTEXT_WINDOW", + "AdaptiveContextCompressor", + "BudgetHealth", "CachedFileContext", + "CompressionMetrics", "ContextCache", "ContextItem", "MemoryStore", "MessageRecord", "SessionRecord", "SharedContextPool", + "TokenBudgetGovernor", + "UsageRecord", + "collapse_repeated_lines", "default_memory_db_path", "estimate_history_tokens", "estimate_message_tokens", "estimate_tokens", "get_default_context_cache", + "strip_ansi_codes", "trim_history_to_budget", ] diff --git a/agentcli/memory/adaptive_compressor.py b/agentcli/memory/adaptive_compressor.py new file mode 100644 index 0000000..62014a4 --- /dev/null +++ b/agentcli/memory/adaptive_compressor.py @@ -0,0 +1,291 @@ +"""Adaptive Context Compression and Tool Output Condensation (Phase 35). + +Progressively compresses tool stdout, large file contents, and older conversation +turns to maximize effective context window utilization without losing critical reasoning. +""" + +from __future__ import annotations + +import hashlib +import logging +import re +from dataclasses import dataclass +from typing import Any + +from ..openrouter_client import ChatMessage +from .budget import ( + DEFAULT_BUDGET_RATIO, + estimate_history_tokens, +) + +logger = logging.getLogger(__name__) + +ANSI_ESCAPE_RE = re.compile(r"\x1b\[[0-9;]*[a-zA-Z]") + + +def strip_ansi_codes(text: str) -> str: + """Remove ANSI escape sequences and terminal formatting codes.""" + return ANSI_ESCAPE_RE.sub("", text) + + +def collapse_repeated_lines(text: str, max_consecutive: int = 2) -> str: + """Collapse sequences of identical repeating lines in logs/stdout.""" + lines = text.splitlines() + if len(lines) <= max_consecutive: + return text + + collapsed: list[str] = [] + prev_line: str | None = None + repeat_count = 0 + + for line in lines: + if line == prev_line: + repeat_count += 1 + if repeat_count < max_consecutive: + collapsed.append(line) + else: + if repeat_count >= max_consecutive: + omitted = repeat_count - max_consecutive + 1 + collapsed.append(f" [... line repeated {omitted} more times ...]") + prev_line = line + repeat_count = 0 + collapsed.append(line) + + if repeat_count >= max_consecutive: + omitted = repeat_count - max_consecutive + 1 + collapsed.append(f" [... line repeated {omitted} more times ...]") + + return "\n".join(collapsed) + + +@dataclass +class CompressionMetrics: + """Statistics for context compression runs.""" + + runs_count: int = 0 + original_tokens: int = 0 + compressed_tokens: int = 0 + + @property + def tokens_saved(self) -> int: + return max(0, self.original_tokens - self.compressed_tokens) + + @property + def compression_ratio(self) -> float: + if self.original_tokens == 0: + return 1.0 + return round(self.compressed_tokens / self.original_tokens, 3) + + def to_dict(self) -> dict[str, Any]: + return { + "runs_count": self.runs_count, + "original_tokens": self.original_tokens, + "compressed_tokens": self.compressed_tokens, + "tokens_saved": self.tokens_saved, + "compression_ratio": self.compression_ratio, + } + + +class AdaptiveContextCompressor: + """Multi-tiered context compression engine with output deduplication.""" + + def __init__( + self, + target_budget_ratio: float = DEFAULT_BUDGET_RATIO, + max_tool_chars: int = 1200, + keep_recent_turns: int = 3, + ) -> None: + self.target_budget_ratio = target_budget_ratio + self.max_tool_chars = max_tool_chars + self.keep_recent_turns = keep_recent_turns + self.metrics = CompressionMetrics() + self._seen_outputs: dict[str, str] = {} + + def compress_tool_output(self, content: str, max_chars: int | None = None) -> str: + """Condense large tool or command output while preserving head and tail context. + + 1. Strips ANSI escape color codes. + 2. Collapses identical repeating lines. + 3. Applies head/tail truncation if exceeding max_chars. + """ + if not content: + return "" + + limit = max_chars or self.max_tool_chars + cleaned = strip_ansi_codes(content) + collapsed = collapse_repeated_lines(cleaned) + + if len(collapsed) <= limit: + return collapsed + + # Deduplication cache check for identical huge outputs + content_hash = hashlib.sha256(collapsed.encode("utf-8")).hexdigest()[:12] + if content_hash in self._seen_outputs and len(collapsed) > limit * 2: + return f"[Tool Output: Identical to previous output (hash: {content_hash}, {len(collapsed):,} chars)]" + self._seen_outputs[content_hash] = collapsed[:100] + + half = limit // 2 + omitted = len(collapsed) - limit + return ( + f"{collapsed[:half]}\n" + f"\n[... {omitted:,} characters omitted by context governor ...]\n\n" + f"{collapsed[-half:]}" + ) + + def compress_tier1(self, history: list[ChatMessage]) -> list[ChatMessage]: + """Tier 1: Prune oversized tool and assistant response outputs.""" + compressed: list[ChatMessage] = [] + for msg in history: + if msg.role in ("tool", "assistant") and len(msg.content or "") > self.max_tool_chars: + compressed.append( + ChatMessage( + role=msg.role, + content=self.compress_tool_output(msg.content or ""), + ) + ) + else: + compressed.append(msg) + return compressed + + def compress_tier2( + self, + history: list[ChatMessage], + keep_recent: int | None = None, + ) -> list[ChatMessage]: + """Tier 2: Synthesize older turn history into structured milestone summaries.""" + if not history: + return [] + + recent_turns = keep_recent if keep_recent is not None else self.keep_recent_turns + system_msg: ChatMessage | None = None + chat_msgs = list(history) + + if chat_msgs and chat_msgs[0].role == "system": + system_msg = chat_msgs.pop(0) + + recent_msg_count = recent_turns * 2 + if len(chat_msgs) <= recent_msg_count: + if system_msg is not None: + return [system_msg, *chat_msgs] + return chat_msgs + + older_msgs = chat_msgs[:-recent_msg_count] + recent_msgs = chat_msgs[-recent_msg_count:] + + summaries: list[str] = [] + for m in older_msgs: + prefix = f"[{m.role.upper()}]" + preview = (m.content or "").strip().replace("\n", " ") + if len(preview) > 140: + preview = f"{preview[:140]}..." + summaries.append(f"{prefix} {preview}") + + summary_text = ( + "[Previous Milestone Context Summary]\n" + + "\n".join(f"- {s}" for s in summaries) + + "\n[End of Milestone Context]" + ) + summary_msg = ChatMessage(role="user", content=summary_text) + + result: list[ChatMessage] = [] + if system_msg is not None: + result.append(system_msg) + result.append(summary_msg) + result.extend(recent_msgs) + return result + + def compress_tier3_emergency( + self, + history: list[ChatMessage], + user_goal: str = "", + touched_files: list[str] | None = None, + ) -> list[ChatMessage]: + """Tier 3: Emergency distillation preserving system instructions, goal, and files.""" + system_msg: ChatMessage | None = None + for msg in history: + if msg.role == "system": + system_msg = msg + break + + files_clause = "" + if touched_files: + files_clause = f"\nTouched Workspace Files: {', '.join(touched_files)}" + + reset_notice = ( + f"[Emergency Context Budget Reset]\n" + f"Active Goal: {user_goal or 'Continue executing task'}{files_clause}\n" + f"Proceeding with latest operational state." + ) + + last_user_or_tool = ( + history[-1] if history else ChatMessage(role="user", content=user_goal) + ) + result: list[ChatMessage] = [] + if system_msg is not None: + result.append(system_msg) + result.append(ChatMessage(role="user", content=reset_notice)) + if last_user_or_tool != system_msg and last_user_or_tool.content != reset_notice: + result.append(last_user_or_tool) + + return result + + def compress( + self, + history: list[ChatMessage], + max_context_tokens: int, + user_goal: str = "", + touched_files: list[str] | None = None, + ) -> list[ChatMessage]: + """Progressively apply compression tiers until history fits the target token budget.""" + if not history: + return [] + + orig_tok = estimate_history_tokens(history) + target_tokens = max(16, int(max_context_tokens * self.target_budget_ratio)) + + if orig_tok <= target_tokens: + return history + + # Pass 1: Prune large tool outputs + t1 = self.compress_tier1(history) + t1_tok = estimate_history_tokens(t1) + if t1_tok <= target_tokens: + self._record_metrics(orig_tok, t1_tok) + return t1 + + # Pass 2: Synthesize older turns + t2 = self.compress_tier2(t1, keep_recent=self.keep_recent_turns) + t2_tok = estimate_history_tokens(t2) + if t2_tok <= target_tokens: + self._record_metrics(orig_tok, t2_tok) + return t2 + + # Pass 3: Tighter Tier 2 with only 1 recent turn + t2_tight = self.compress_tier2(t1, keep_recent=1) + t2_tight_tok = estimate_history_tokens(t2_tight) + if t2_tight_tok <= target_tokens: + self._record_metrics(orig_tok, t2_tight_tok) + return t2_tight + + # Pass 4: Tier 3 Emergency Context Reset + t3 = self.compress_tier3_emergency( + history=history, + user_goal=user_goal, + touched_files=touched_files, + ) + t3_tok = estimate_history_tokens(t3) + self._record_metrics(orig_tok, t3_tok) + return t3 + + def _record_metrics(self, original_tokens: int, compressed_tokens: int) -> None: + self.metrics.runs_count += 1 + self.metrics.original_tokens += original_tokens + self.metrics.compressed_tokens += compressed_tokens + + +__all__ = [ + "AdaptiveContextCompressor", + "CompressionMetrics", + "collapse_repeated_lines", + "strip_ansi_codes", +] diff --git a/agentcli/memory/governor.py b/agentcli/memory/governor.py new file mode 100644 index 0000000..79d7467 --- /dev/null +++ b/agentcli/memory/governor.py @@ -0,0 +1,365 @@ +"""Token Budget Governor and Real-Time Accounting (Phase 35). + +Provides dynamic token and cost tracking, soft warning thresholds, hard ceilings, +spend velocity calculations ($/hr, tokens/min), and per-agent usage breakdowns. +""" + +from __future__ import annotations + +import logging +import time +from dataclasses import dataclass +from typing import Any + +from .budget import calculate_cost, estimate_tokens + +logger = logging.getLogger(__name__) + + +@dataclass +class UsageRecord: + """Record of a single model invocation token usage and cost.""" + + timestamp: float + model: str + prompt_tokens: int + completion_tokens: int + cached_tokens: int = 0 + cost_usd: float = 0.0 + agent_type: str = "main" + + @property + def total_tokens(self) -> int: + return self.prompt_tokens + self.completion_tokens + + +@dataclass +class BudgetHealth: + """Snapshot of budget health status and velocity metrics.""" + + status: str # "ok", "warning", "exceeded" + used_cost_usd: float + max_cost_usd: float | None + used_tokens: int + max_tokens: int | None + utilization_pct: float + cost_per_hour: float + tokens_per_minute: float + is_warning: bool + is_exceeded: bool + warning_ratio: float + + def to_dict(self) -> dict[str, Any]: + return { + "status": self.status, + "used_cost_usd": round(self.used_cost_usd, 6), + "max_cost_usd": self.max_cost_usd, + "used_tokens": self.used_tokens, + "max_tokens": self.max_tokens, + "utilization_pct": round(self.utilization_pct, 2), + "cost_per_hour": round(self.cost_per_hour, 6), + "tokens_per_minute": round(self.tokens_per_minute, 2), + "is_warning": self.is_warning, + "is_exceeded": self.is_exceeded, + "warning_ratio": self.warning_ratio, + } + + +class TokenBudgetGovernor: + """Real-time token usage governor with budget enforcement and velocity metrics.""" + + def __init__( + self, + max_cost_usd: float | None = None, + max_tokens: int | None = None, + warning_ratio: float = 0.80, + ) -> None: + self.max_cost_usd = max_cost_usd + self.max_tokens = max_tokens + self.warning_ratio = max(0.1, min(0.99, warning_ratio)) + self._records: list[UsageRecord] = [] + self._start_time = time.time() + + @property + def total_cost_usd(self) -> float: + """Total cumulative spend across all recorded invocations in USD.""" + return sum(r.cost_usd for r in self._records) + + @property + def total_tokens(self) -> int: + """Total token count (prompt + completion) across all invocations.""" + return sum(r.total_tokens for r in self._records) + + @property + def prompt_tokens(self) -> int: + """Total prompt/input tokens consumed.""" + return sum(r.prompt_tokens for r in self._records) + + @property + def completion_tokens(self) -> int: + """Total completion/generated tokens consumed.""" + return sum(r.completion_tokens for r in self._records) + + @property + def cached_tokens(self) -> int: + """Total cached or deduplicated tokens saved.""" + return sum(r.cached_tokens for r in self._records) + + @property + def records(self) -> list[UsageRecord]: + return list(self._records) + + def record_usage( + self, + model: str, + prompt_tokens: int, + completion_tokens: int, + cached_tokens: int = 0, + agent_type: str = "main", + ) -> float: + """Record model invocation token usage and calculate USD cost. + + Args: + model: Model identifier string. + prompt_tokens: Number of prompt/input tokens. + completion_tokens: Number of output/completion tokens. + cached_tokens: Number of prompt tokens read from cache. + agent_type: Subagent type or 'main'. + + Returns: + Calculated cost in USD for this invocation. + """ + cost = calculate_cost(model, prompt_tokens, completion_tokens) + record = UsageRecord( + timestamp=time.time(), + model=model, + prompt_tokens=max(0, prompt_tokens), + completion_tokens=max(0, completion_tokens), + cached_tokens=max(0, cached_tokens), + cost_usd=cost, + agent_type=agent_type or "main", + ) + self._records.append(record) + + if self.is_exceeded(): + logger.warning( + "Budget ceiling exceeded! Total: $%.4f (Limit: $%.4f)", + self.total_cost_usd, + self.max_cost_usd or 0.0, + ) + elif self.is_warning(): + logger.info( + "Approaching budget ceiling (%.1f%% utilized). Total: $%.4f", + self.get_utilization_pct(), + self.total_cost_usd, + ) + + return cost + + def estimate_and_record_text( + self, + model: str, + prompt_text: str, + completion_text: str, + cached_tokens: int = 0, + agent_type: str = "main", + ) -> float: + """Estimate token counts from raw strings using character heuristics and record usage.""" + p_tok = estimate_tokens(prompt_text) + c_tok = estimate_tokens(completion_text) + return self.record_usage( + model=model, + prompt_tokens=p_tok, + completion_tokens=c_tok, + cached_tokens=cached_tokens, + agent_type=agent_type, + ) + + def is_exceeded(self) -> bool: + """Check if cumulative cost or tokens have reached or exceeded hard ceilings.""" + if self.max_cost_usd is not None and self.total_cost_usd >= self.max_cost_usd: + return True + return bool(self.max_tokens is not None and self.total_tokens >= self.max_tokens) + + def is_warning(self) -> bool: + """Check if cumulative cost or tokens have reached soft warning threshold.""" + if self.is_exceeded(): + return True + if ( + self.max_cost_usd is not None + and self.total_cost_usd >= self.max_cost_usd * self.warning_ratio + ): + return True + return bool( + self.max_tokens is not None + and self.total_tokens >= self.max_tokens * self.warning_ratio + ) + + + def get_utilization_pct(self) -> float: + """Get percentage utilization of the most constrained budget ceiling.""" + percentages: list[float] = [] + if self.max_cost_usd is not None and self.max_cost_usd > 0: + percentages.append((self.total_cost_usd / self.max_cost_usd) * 100.0) + if self.max_tokens is not None and self.max_tokens > 0: + percentages.append((self.total_tokens / self.max_tokens) * 100.0) + return max(percentages) if percentages else 0.0 + + def cost_per_hour(self) -> float: + """Calculate spend velocity in USD per hour.""" + if not self._records: + return 0.0 + elapsed_seconds = max(1.0, time.time() - self._start_time) + return (self.total_cost_usd / elapsed_seconds) * 3600.0 + + def tokens_per_minute(self) -> float: + """Calculate token throughput in tokens per minute.""" + if not self._records: + return 0.0 + elapsed_seconds = max(1.0, time.time() - self._start_time) + return (self.total_tokens / elapsed_seconds) * 60.0 + + def get_agent_breakdown(self) -> dict[str, dict[str, Any]]: + """Return token count, cost, and invocation count grouped by agent type.""" + breakdown: dict[str, dict[str, Any]] = {} + for r in self._records: + agent = r.agent_type or "main" + if agent not in breakdown: + breakdown[agent] = { + "calls": 0, + "prompt_tokens": 0, + "completion_tokens": 0, + "total_tokens": 0, + "cached_tokens": 0, + "cost_usd": 0.0, + } + breakdown[agent]["calls"] += 1 + breakdown[agent]["prompt_tokens"] += r.prompt_tokens + breakdown[agent]["completion_tokens"] += r.completion_tokens + breakdown[agent]["total_tokens"] += r.total_tokens + breakdown[agent]["cached_tokens"] += r.cached_tokens + breakdown[agent]["cost_usd"] += r.cost_usd + + # Add percentage shares + total_cost = self.total_cost_usd + for data in breakdown.values(): + data["cost_usd"] = round(data["cost_usd"], 6) + data["cost_share_pct"] = ( + round((data["cost_usd"] / total_cost) * 100.0, 1) if total_cost > 0 else 0.0 + ) + + return breakdown + + def get_model_breakdown(self) -> dict[str, dict[str, Any]]: + """Return token count, cost, and invocation count grouped by model.""" + breakdown: dict[str, dict[str, Any]] = {} + for r in self._records: + m = r.model or "unknown" + if m not in breakdown: + breakdown[m] = { + "calls": 0, + "prompt_tokens": 0, + "completion_tokens": 0, + "total_tokens": 0, + "cost_usd": 0.0, + } + breakdown[m]["calls"] += 1 + breakdown[m]["prompt_tokens"] += r.prompt_tokens + breakdown[m]["completion_tokens"] += r.completion_tokens + breakdown[m]["total_tokens"] += r.total_tokens + breakdown[m]["cost_usd"] += r.cost_usd + + for data in breakdown.values(): + data["cost_usd"] = round(data["cost_usd"], 6) + + return breakdown + + def check_health(self) -> BudgetHealth: + """Generate a complete budget health snapshot.""" + if self.is_exceeded(): + status = "exceeded" + elif self.is_warning(): + status = "warning" + else: + status = "ok" + + return BudgetHealth( + status=status, + used_cost_usd=self.total_cost_usd, + max_cost_usd=self.max_cost_usd, + used_tokens=self.total_tokens, + max_tokens=self.max_tokens, + utilization_pct=self.get_utilization_pct(), + cost_per_hour=self.cost_per_hour(), + tokens_per_minute=self.tokens_per_minute(), + is_warning=self.is_warning(), + is_exceeded=self.is_exceeded(), + warning_ratio=self.warning_ratio, + ) + + def set_budget( + self, + max_cost_usd: float | None = None, + max_tokens: int | None = None, + warning_ratio: float | None = None, + ) -> None: + """Dynamically update budget limits and warning ratios.""" + if max_cost_usd is not None: + self.max_cost_usd = max(0.0, float(max_cost_usd)) if max_cost_usd > 0 else None + if max_tokens is not None: + self.max_tokens = max(0, int(max_tokens)) if max_tokens > 0 else None + if warning_ratio is not None: + self.warning_ratio = max(0.1, min(0.99, float(warning_ratio))) + + def reset(self) -> None: + """Reset all usage records and reset session timer.""" + self._records.clear() + self._start_time = time.time() + + def format_summary(self) -> str: + """Generate a human-readable CLI summary table of cost, tokens, and velocity.""" + health = self.check_health() + lines: list[str] = [ + "================== Token & Cost Budget ==================", + f" Cumulative Cost : ${health.used_cost_usd:.4f}" + + (f" / ${health.max_cost_usd:.4f}" if health.max_cost_usd else " (No limit)"), + f" Total Tokens : {health.used_tokens:,}" + + (f" / {health.max_tokens:,}" if health.max_tokens else ""), + f" - Prompt : {self.prompt_tokens:,}", + f" - Completion : {self.completion_tokens:,}", + f" - Cached/Saved: {self.cached_tokens:,}", + f" Utilization : {health.utilization_pct:.1f}%", + f" Spend Velocity : ${health.cost_per_hour:.4f}/hr ({health.tokens_per_minute:.0f} tok/min)", + f" Status : {health.status.upper()}", + "---------------------------------------------------------", + ] + + # Breakdown by Agent + agent_bd = self.get_agent_breakdown() + if agent_bd: + lines.append("By SubAgent:") + for agent, stats in agent_bd.items(): + lines.append( + f" * {agent:<14} : {stats['calls']:>2} calls | {stats['total_tokens']:>7,} tok | ${stats['cost_usd']:.4f} ({stats['cost_share_pct']}%)" + ) + lines.append("---------------------------------------------------------") + + # Breakdown by Model + model_bd = self.get_model_breakdown() + if model_bd: + lines.append("By Model:") + for model, stats in model_bd.items(): + short_m = model.split("/")[-1] if "/" in model else model + lines.append( + f" * {short_m:<20} : {stats['calls']:>2} calls | {stats['total_tokens']:>7,} tok | ${stats['cost_usd']:.4f}" + ) + lines.append("=========================================================") + + return "\n".join(lines) + + +__all__ = [ + "BudgetHealth", + "TokenBudgetGovernor", + "UsageRecord", +] diff --git a/agentcli/session.py b/agentcli/session.py index cda41b0..c507a8e 100644 --- a/agentcli/session.py +++ b/agentcli/session.py @@ -15,7 +15,9 @@ from .config import Config from .files import load_agents_md from .mcp.manager import MCPClientManager +from .memory.adaptive_compressor import AdaptiveContextCompressor from .memory.budget import DEFAULT_CONTEXT_WINDOW, estimate_tokens, trim_history_to_budget +from .memory.governor import TokenBudgetGovernor from .memory.store import MemoryStore from .openrouter_client import ( ChatMessage, @@ -109,6 +111,11 @@ def __init__( self.history.insert(0, ChatMessage(role="system", content=agents_context)) self.cumulative_cost_usd: float = 0.0 + max_cost = getattr(config.routing, "max_cost_usd", None) + self.governor: TokenBudgetGovernor = TokenBudgetGovernor(max_cost_usd=max_cost) + self.compressor: AdaptiveContextCompressor = AdaptiveContextCompressor( + target_budget_ratio=config.memory.budget_ratio + ) self.registry: ModelRegistry | None = None self.router: Router | None = None self.mcp_manager: MCPClientManager = MCPClientManager(config=self.config) @@ -133,18 +140,35 @@ async def initialize_mcp(self) -> None: if self.config.mcp_servers: await self.mcp_manager.initialize() - def record_cost(self, model: str, prompt_tokens: int, completion_tokens: int) -> float: + def record_cost( + self, + model: str, + prompt_tokens: int, + completion_tokens: int, + cached_tokens: int = 0, + agent_type: str = "main", + ) -> float: """Calculate and accumulate the USD cost for a model invocation.""" - from .memory.budget import calculate_cost - - cost = calculate_cost(model, prompt_tokens, completion_tokens) - self.cumulative_cost_usd += cost + cost = self.governor.record_usage( + model=model, + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + cached_tokens=cached_tokens, + agent_type=agent_type, + ) + self.cumulative_cost_usd = self.governor.total_cost_usd return cost def is_budget_exceeded(self) -> bool: """Check if session cumulative cost has reached or exceeded max_cost_usd.""" - max_cost = getattr(self.config.routing, "max_cost_usd", None) - return bool(max_cost is not None and self.cumulative_cost_usd >= max_cost) + max_cost = getattr(self.config.routing, "max_cost_usd", None) or self.governor.max_cost_usd + if ( + max_cost is not None + and max_cost > 0 + and max(self.cumulative_cost_usd, self.governor.total_cost_usd) >= max_cost + ): + return True + return self.governor.is_exceeded() def close(self) -> None: store = getattr(self, "memory_store", None) @@ -173,18 +197,20 @@ def _resolve_context_window(self, model_id: str | None) -> int: return DEFAULT_CONTEXT_WINDOW def _trim_history(self, max_context_tokens: int | None = None) -> list[ChatMessage]: - """Trim conversation history using dynamic token budget bounded by history_turns.""" + """Trim conversation history using dynamic token budget and adaptive compression.""" target_window = ( max_context_tokens if max_context_tokens is not None else self._resolve_context_window(self.forced_model) ) - return trim_history_to_budget( + trimmed = trim_history_to_budget( self.history, max_context_tokens=target_window, max_turns=self.config.app.history_turns, budget_ratio=self.config.memory.budget_ratio, ) + return self.compressor.compress(trimmed, max_context_tokens=target_window) + def add_user_message(self, content: str, token_count: int | None = None) -> None: self.history.append(ChatMessage(role="user", content=content)) diff --git a/agentcli/ui/prompt.py b/agentcli/ui/prompt.py index 9700ab2..0f7e31c 100644 --- a/agentcli/ui/prompt.py +++ b/agentcli/ui/prompt.py @@ -170,12 +170,16 @@ def get_completions(self, document: Document, complete_event: CompleteEvent) -> if text.startswith(("/budget ", "\\budget ")): arg = text.split(maxsplit=1)[1] if len(text.split(maxsplit=1)) > 1 else "" arg_lower = arg.lower() - tier_options = [ + budget_options = [ + ("status", "[ACTION] Show budget usage and velocity summary"), + ("set", "[ACTION] Set USD budget ceiling (/budget set )"), + ("max-tokens", "[ACTION] Set token ceiling (/budget max-tokens )"), + ("reset", "[ACTION] Reset session cost and token counters"), ("low", "[TIER] Free models only"), ("medium", "[TIER] High-efficiency & free models"), ("high", "[TIER] Frontier reasoning & coding models"), ] - for opt, desc in tier_options: + for opt, desc in budget_options: if opt.startswith(arg_lower): yield Completion(opt, start_position=-len(arg), display_meta=desc) return diff --git a/agentcli/ui/tui_app.py b/agentcli/ui/tui_app.py index fda4188..eebaf0a 100644 --- a/agentcli/ui/tui_app.py +++ b/agentcli/ui/tui_app.py @@ -641,6 +641,51 @@ async def _handle_slash_command(self, text: str, timestamp: str, event: KeyPress self.add_message("system", "Worktree manager not available in active session.", timestamp) return + if cmd in ("/budget", "/cost", "/tokens"): + if self.session and hasattr(self.session, "governor"): + parts = text.split(maxsplit=2) + subcmd = parts[1].lower() if len(parts) > 1 else "status" + val = parts[2].strip() if len(parts) > 2 else "" + + if cmd in ("/cost", "/tokens") or subcmd in ("status", "info", "show") or len(parts) == 1: + self.add_message("system", self.session.governor.format_summary(), timestamp) + elif subcmd in ("set", "limit", "max-cost"): + if not val: + self.add_message("system", "Usage: /budget set ", timestamp) + else: + try: + f_val = float(val.lstrip("$")) + self.session.governor.set_budget(max_cost_usd=f_val) + self.add_message("system", f"Session budget ceiling set to: ${f_val:.4f} USD", timestamp) + except ValueError: + self.add_message("system", f"Invalid budget amount: '{val}'.", timestamp) + elif subcmd in ("tokens", "max-tokens"): + if not val: + self.add_message("system", "Usage: /budget max-tokens ", timestamp) + else: + try: + t_val = int(val.replace(",", "")) + self.session.governor.set_budget(max_tokens=t_val) + self.add_message("system", f"Session token budget ceiling set to: {t_val:,} tokens", timestamp) + except ValueError: + self.add_message("system", f"Invalid token count: '{val}'.", timestamp) + elif subcmd in ("reset", "clear"): + self.session.governor.reset() + self.session.cumulative_cost_usd = 0.0 + self.add_message("system", "Budget and cost counters reset for current session.", timestamp) + elif subcmd in ("low", "medium", "high", "tier"): + tier_name = val if subcmd == "tier" else subcmd + if tier_name in ("low", "medium", "high"): + self.session.config.routing.budget_tier = tier_name + self.add_message("system", f"Budget routing tier switched to: [{tier_name.upper()}]", timestamp) + else: + self.add_message("system", "Usage: /budget tier [low | medium | high]", timestamp) + else: + self.add_message("system", "Usage: /budget [status | set | max-tokens | tier | reset]", timestamp) + else: + self.add_message("system", "Budget governor not available in active session.", timestamp) + return + if cmd == "/goal": parts = text.split(maxsplit=1) diff --git a/tests/test_phase35_token_governor.py b/tests/test_phase35_token_governor.py new file mode 100644 index 0000000..acf1c11 --- /dev/null +++ b/tests/test_phase35_token_governor.py @@ -0,0 +1,341 @@ +from __future__ import annotations + +import pytest + +from agentcli.config import Config +from agentcli.memory.adaptive_compressor import ( + AdaptiveContextCompressor, + collapse_repeated_lines, + strip_ansi_codes, +) +from agentcli.memory.governor import BudgetHealth, TokenBudgetGovernor +from agentcli.openrouter_client import ChatMessage +from agentcli.session import AgentSession + + +def test_governor_record_usage_and_costs() -> None: + """Test TokenBudgetGovernor records token usage, cached tokens, and calculates USD cost.""" + governor = TokenBudgetGovernor(max_cost_usd=1.00) + + # 1. Free model should cost $0.00 + cost1 = governor.record_usage( + model="google/gemma-4-31b-it:free", + prompt_tokens=1000, + completion_tokens=500, + cached_tokens=200, + agent_type="planner", + ) + assert cost1 == 0.0 + assert governor.total_tokens == 1500 + assert governor.prompt_tokens == 1000 + assert governor.completion_tokens == 500 + assert governor.cached_tokens == 200 + assert governor.total_cost_usd == 0.0 + + # 2. Paid model should accumulate cost + cost2 = governor.record_usage( + model="openai/gpt-4o", + prompt_tokens=10_000, # $2.50 / 1M = $0.025 + completion_tokens=2_000, # $10.00 / 1M = $0.020 -> total $0.045 + agent_type="code_analyzer", + ) + assert cost2 == pytest.approx(0.045, rel=1e-3) + assert governor.total_cost_usd == pytest.approx(0.045, rel=1e-3) + assert governor.total_tokens == 13_500 + assert len(governor.records) == 2 + + +def test_governor_soft_warning_and_hard_ceiling() -> None: + """Test soft warning threshold (80%) and hard budget ceiling (100%).""" + governor = TokenBudgetGovernor(max_cost_usd=0.10, warning_ratio=0.80) + + assert governor.is_warning() is False + assert governor.is_exceeded() is False + + # Incur cost of $0.05 (50% - OK) + governor.record_usage( + model="anthropic/claude-3.5-sonnet", + prompt_tokens=10_000, # 10k * $3.00/1M = $0.03 + completion_tokens=1_333, # 1.33k * $15.00/1M = $0.02 + ) + assert governor.is_warning() is False + assert governor.is_exceeded() is False + assert 45.0 <= governor.get_utilization_pct() <= 55.0 + + # Incur additional cost to cross 80% ($0.085 total) + governor.record_usage( + model="anthropic/claude-3.5-sonnet", + prompt_tokens=10_000, + completion_tokens=1_000, + ) + assert governor.is_warning() is True + assert governor.is_exceeded() is False + + # Incur additional cost to cross 100% ($0.12 total) + governor.record_usage( + model="anthropic/claude-3.5-sonnet", + prompt_tokens=10_000, + completion_tokens=2_000, + ) + assert governor.is_warning() is True + assert governor.is_exceeded() is True + assert governor.get_utilization_pct() >= 100.0 + + +def test_governor_token_limit_and_health() -> None: + """Test max_tokens ceiling and check_health snapshot generation.""" + governor = TokenBudgetGovernor(max_tokens=5000, warning_ratio=0.75) + + governor.record_usage( + model="google/gemma-4-31b-it:free", + prompt_tokens=2000, + completion_tokens=1000, + ) + health = governor.check_health() + assert isinstance(health, BudgetHealth) + assert health.status == "ok" + assert health.used_tokens == 3000 + assert health.max_tokens == 5000 + assert health.utilization_pct == 60.0 + + # Cross 75% soft warning + governor.record_usage( + model="google/gemma-4-31b-it:free", + prompt_tokens=1000, + completion_tokens=500, + ) + assert governor.is_warning() is True + assert governor.check_health().status == "warning" + + # Cross 100% token limit + governor.record_usage( + model="google/gemma-4-31b-it:free", + prompt_tokens=1000, + completion_tokens=1000, + ) + assert governor.is_exceeded() is True + assert governor.check_health().status == "exceeded" + + # Test health dict serialization + h_dict = health.to_dict() + assert "used_cost_usd" in h_dict + assert "utilization_pct" in h_dict + assert "cost_per_hour" in h_dict + + +def test_governor_velocity_and_breakdowns() -> None: + """Test spend velocity, subagent breakdown, and model breakdown.""" + governor = TokenBudgetGovernor(max_cost_usd=1.00) + + governor.record_usage( + model="openai/gpt-4o-mini", + prompt_tokens=5000, + completion_tokens=1000, + agent_type="planner", + ) + governor.record_usage( + model="openai/gpt-4o-mini", + prompt_tokens=3000, + completion_tokens=500, + agent_type="code_analyzer", + ) + governor.record_usage( + model="google/gemma-4-31b-it:free", + prompt_tokens=4000, + completion_tokens=200, + agent_type="planner", + ) + + # Agent breakdown + agent_bd = governor.get_agent_breakdown() + assert "planner" in agent_bd + assert "code_analyzer" in agent_bd + assert agent_bd["planner"]["calls"] == 2 + assert agent_bd["planner"]["total_tokens"] == 10_200 + assert agent_bd["code_analyzer"]["calls"] == 1 + + # Model breakdown + model_bd = governor.get_model_breakdown() + assert "openai/gpt-4o-mini" in model_bd + assert "google/gemma-4-31b-it:free" in model_bd + assert model_bd["openai/gpt-4o-mini"]["calls"] == 2 + + # Spend velocity calculations + vel_cost = governor.cost_per_hour() + vel_tok = governor.tokens_per_minute() + assert vel_cost >= 0.0 + assert vel_tok >= 0.0 + + # Format summary table + summary_text = governor.format_summary() + assert "Token & Cost Budget" in summary_text + assert "By SubAgent:" in summary_text + assert "By Model:" in summary_text + + +def test_governor_set_budget_and_reset() -> None: + """Test dynamic runtime budget adjustment and reset.""" + governor = TokenBudgetGovernor(max_cost_usd=0.50, max_tokens=10000) + governor.record_usage("openai/gpt-4o", 1000, 500) + assert governor.total_tokens == 1500 + assert len(governor.records) == 1 + + # Update budget limits + governor.set_budget(max_cost_usd=2.00, max_tokens=50000, warning_ratio=0.85) + assert governor.max_cost_usd == 2.00 + assert governor.max_tokens == 50000 + assert governor.warning_ratio == 0.85 + + # Reset + governor.reset() + assert governor.total_tokens == 0 + assert governor.total_cost_usd == 0.0 + assert len(governor.records) == 0 + + +def test_governor_estimate_and_record_text() -> None: + """Test estimating tokens from raw strings and recording.""" + governor = TokenBudgetGovernor() + prompt = "Review this authentication module and check for token leaks." + completion = "The code looks solid. No leaks found." + + cost = governor.estimate_and_record_text( + model="openai/gpt-4o-mini", + prompt_text=prompt, + completion_text=completion, + cached_tokens=10, + agent_type="reviewer", + ) + assert cost >= 0.0 + assert governor.prompt_tokens > 0 + assert governor.completion_tokens > 0 + assert governor.cached_tokens == 10 + + +def test_strip_ansi_and_collapse_lines() -> None: + """Test ANSI strip and repeated line collapsing utilities.""" + # 1. ANSI strip + colored = "\x1b[31mError:\x1b[0m File \x1b[1mnot found\x1b[0m" + assert strip_ansi_codes(colored) == "Error: File not found" + + # 2. Collapse repeating lines + log_text = "Checking...\n" + ("Downloading package xyz\n" * 10) + "Done." + collapsed = collapse_repeated_lines(log_text, max_consecutive=2) + assert "repeated 8 more times" in collapsed + assert "Done." in collapsed + + +def test_adaptive_compressor_tool_output() -> None: + """Test tool output truncation and hash deduplication.""" + compressor = AdaptiveContextCompressor(max_tool_chars=200) + + # Short output remains unchanged + short_text = "All 10 tests passed." + assert compressor.compress_tool_output(short_text) == short_text + + # Long output gets truncated (using unique lines so repeated line collapse doesn't shrink it first) + unique_lines = "\n".join(f"data record row #{i} from server output stream" for i in range(40)) + long_text = f"START_BLOCK\n{unique_lines}\nEND_BLOCK" + pruned = compressor.compress_tool_output(long_text) + assert "START_BLOCK" in pruned + assert "END_BLOCK" in pruned + assert "characters omitted by context governor" in pruned + assert len(pruned) < len(long_text) + + # Identical huge output is deduplicated + huge_text = "UNIQUE_HUGE_STAMP_" + ("x" * 1000) + compressor.compress_tool_output(huge_text) + out2 = compressor.compress_tool_output(huge_text) + assert "Identical to previous output" in out2 + + + + +def test_adaptive_compressor_progressive_tiers() -> None: + """Test progressive compression through Tier 1, Tier 2, and Tier 3 emergency.""" + compressor = AdaptiveContextCompressor( + target_budget_ratio=0.75, + max_tool_chars=100, + keep_recent_turns=1, + ) + + sys_msg = ChatMessage(role="system", content="You are a developer assistant.") + u1 = ChatMessage(role="user", content="Task 1: Read files") + a1 = ChatMessage(role="assistant", content="Reading file content:\n" + ("data line\n" * 30)) + u2 = ChatMessage(role="user", content="Task 2: Refactor auth") + a2 = ChatMessage(role="assistant", content="Refactored auth module successfully.") + u3 = ChatMessage(role="user", content="Task 3: Run test suite") + + history = [sys_msg, u1, a1, u2, a2, u3] + + # Compress history to fit within small token window (e.g. 100 tokens) + compressed = compressor.compress( + history=history, + max_context_tokens=150, + user_goal="Refactor auth and run tests", + touched_files=["auth.py", "test_auth.py"], + ) + + assert len(compressed) >= 2 + assert compressed[0].role == "system" + # Ensure metrics were tracked + assert compressor.metrics.runs_count >= 1 + assert compressor.metrics.tokens_saved >= 0 + assert 0.0 < compressor.metrics.compression_ratio <= 1.0 + + # Emergency compression + emergency = compressor.compress_tier3_emergency( + history=history, + user_goal="Critical emergency fix", + touched_files=["main.py"], + ) + assert emergency[0].role == "system" + assert any("[Emergency Context Budget Reset]" in (m.content or "") for m in emergency) + + +@pytest.mark.asyncio +async def test_session_governor_integration(monkeypatch: pytest.MonkeyPatch) -> None: + """Test AgentSession integrates with TokenBudgetGovernor for cost and limits.""" + monkeypatch.setenv("OPENROUTER_API_KEY", "test-key") + config = Config() + config.routing.max_cost_usd = 0.05 + + session = AgentSession(config=config, forced_model="google/gemma-4-31b-it:free") + assert session.governor is not None + assert session.compressor is not None + assert session.is_budget_exceeded() is False + + # Record cost through session + cost = session.record_cost("openai/gpt-4o-mini", prompt_tokens=1000, completion_tokens=500) + assert cost > 0.0 + assert session.cumulative_cost_usd == session.governor.total_cost_usd + + # Trigger budget ceiling + session.governor.set_budget(max_cost_usd=0.0001) + assert session.is_budget_exceeded() is True + + # When budget exceeded, step returns warning message + reply = await session.step("Should not execute") + assert "Session cost ceiling reached" in reply + + await session.aclose() + + +def test_slash_completer_budget_completions() -> None: + """Test SlashAndFileCompleter offers new /budget subcommands.""" + from prompt_toolkit.completion import CompleteEvent + from prompt_toolkit.document import Document + + from agentcli.ui.prompt import SlashAndFileCompleter + + completer = SlashAndFileCompleter() + event = CompleteEvent() + + comps = [c.text for c in completer.get_completions(Document("/budget "), event)] + assert "status" in comps + assert "set" in comps + assert "max-tokens" in comps + assert "reset" in comps + assert "low" in comps + assert "medium" in comps + assert "high" in comps From 8ac911c89ff2478df86630a126251c5f1974e2a1 Mon Sep 17 00:00:00 2001 From: De-pitcher Date: Mon, 7 Sep 2026 06:06:17 +0100 Subject: [PATCH 2/4] test(phase35): expand governor, adaptive compressor, and subagents coverage --- agentcli/ui/tui_app.py | 79 +++++++-- tests/test_coverage_boost.py | 93 ++++++++++ tests/test_phase35_token_governor.py | 244 +++++++++++++++++++++++++++ 3 files changed, 399 insertions(+), 17 deletions(-) diff --git a/agentcli/ui/tui_app.py b/agentcli/ui/tui_app.py index eebaf0a..c086c97 100644 --- a/agentcli/ui/tui_app.py +++ b/agentcli/ui/tui_app.py @@ -269,25 +269,69 @@ async def _handle_slash_command(self, text: str, timestamp: str, event: KeyPress self._app.invalidate() return - if cmd == "/budget": - parts = text.split(maxsplit=1) - if len(parts) == 1: - self.add_message( - "system", f"Current budget tier: {self.config.routing.budget_tier}", timestamp - ) - else: - tier = parts[1].strip().lower() - if tier in {"low", "medium", "high"}: - self.config.routing.budget_tier = tier + if cmd in {"/budget", "\\budget"}: + parts = text.split(maxsplit=2) + subcmd = parts[1].lower() if len(parts) > 1 else "" + val = parts[2].strip() if len(parts) > 2 else "" + + if not subcmd or subcmd in ("status", "info", "show"): + curr = self.config.routing.budget_tier + msg = f"Current budget tier: {curr}" + if self.session and hasattr(self.session, "governor"): + msg += "\n" + self.session.governor.format_summary() + self.add_message("system", msg, timestamp) + elif subcmd in ("low", "medium", "high", "tier"): + tier_name = val if subcmd == "tier" else subcmd + if tier_name in ("low", "medium", "high"): + self.config.routing.budget_tier = tier_name if self.session and self.session.router is not None: - self.session.router._budget_tier = tier - self.add_message("system", f"Budget tier updated to: {tier}", timestamp) + self.session.router._budget_tier = tier_name + self.add_message("system", f"Budget tier updated to: {tier_name}", timestamp) else: self.add_message( "system", - f"Invalid budget tier '{tier}'. Choose from: low, medium, high", + f"Invalid budget tier '{tier_name}'. Choose from: low, medium, high", timestamp, ) + elif subcmd in ("set", "limit", "max-cost"): + if not val: + self.add_message("system", "Usage: /budget set ", timestamp) + else: + try: + cost_limit = float(val.lstrip("$")) + if self.session and hasattr(self.session, "governor"): + self.session.governor.set_budget(max_cost_usd=cost_limit) + self.config.routing.max_cost_usd = cost_limit + self.state.budget_limit_usd = cost_limit + self.add_message( + "system", f"Session budget ceiling set to: ${cost_limit:.4f} USD", timestamp + ) + except ValueError: + self.add_message( + "system", f"Invalid budget amount: '{val}'. Must be a positive number.", timestamp + ) + elif subcmd in ("tokens", "max-tokens"): + if not val: + self.add_message("system", "Usage: /budget max-tokens ", timestamp) + else: + try: + t_count = int(val.replace(",", "")) + if self.session and hasattr(self.session, "governor"): + self.session.governor.set_budget(max_tokens=t_count) + self.add_message( + "system", f"Session token budget ceiling set to: {t_count:,} tokens", timestamp + ) + except ValueError: + self.add_message("system", f"Invalid token count: '{val}'.", timestamp) + elif subcmd in ("reset", "clear"): + if self.session and hasattr(self.session, "governor"): + self.session.governor.reset() + self.session.cumulative_cost_usd = 0.0 + self.add_message("system", "Budget and cost counters reset for current session.", timestamp) + else: + self.add_message( + "system", f"Invalid budget subcommand '{subcmd}'. Choose from: low, medium, high, status, set, max-tokens, reset", timestamp + ) return if cmd in {"/models", "/model"}: @@ -357,13 +401,14 @@ async def _handle_slash_command(self, text: str, timestamp: str, event: KeyPress - self.state.completion_tokens, cost_usd=cost - self.state.cost_usd, ) - self.add_message( - "system", + msg_text = ( f"Token Usage: {stats['total_tokens']} total " f"({stats['user_tokens']} prompt, {stats['assistant_tokens']} completion) | " - f"Estimated Cost: ${cost:.6f} USD", - timestamp, + f"Estimated Cost: ${cost:.6f} USD" ) + if hasattr(self.session, "governor"): + msg_text += "\n" + self.session.governor.format_summary() + self.add_message("system", msg_text, timestamp) return if cmd in {"/clear", "/cls"}: diff --git a/tests/test_coverage_boost.py b/tests/test_coverage_boost.py index 6f8d4e2..5fe79fc 100644 --- a/tests/test_coverage_boost.py +++ b/tests/test_coverage_boost.py @@ -470,5 +470,98 @@ async def mock_stream(*args, **kwargs): assert "LLM analysis failed" in (res_err.error or "") +@pytest.mark.asyncio +async def test_workspace_agent_full_coverage(tmp_path: Path) -> None: + """Test WorkspaceAgent branches including list_tree, git_branch, and grep.""" + from agentcli.subagents.base import SubAgentTask, SubAgentType + from agentcli.subagents.workspace import WorkspaceAgent + + agent = WorkspaceAgent() + d1 = tmp_path / "src" + d1.mkdir() + (d1 / "app.py").write_text("print('hello world')\n", encoding="utf-8") + (tmp_path / "README.md").write_text("# Test Repo\n", encoding="utf-8") + + # 1. list_tree + task_tree = SubAgentTask( + agent_type=SubAgentType.WORKSPACE, + payload={"operation": "list_tree", "path": str(tmp_path), "max_depth": 3}, + ) + res_tree = await agent.run(task_tree) + assert res_tree.success is True + assert any("app.py" in line for line in res_tree.output["tree"]) + + # 2. search_code + task_search = SubAgentTask( + agent_type=SubAgentType.WORKSPACE, + payload={"operation": "search_code", "query": "hello", "path": str(tmp_path)}, + ) + res_search = await agent.run(task_search) + assert res_search.success is True + assert res_search.output["total_matches"] >= 1 + + # 3. git_branch list / create + task_git_list = SubAgentTask( + agent_type=SubAgentType.WORKSPACE, + payload={"operation": "git_branch", "action": "list", "path": str(tmp_path)}, + ) + res_git = await agent.run(task_git_list) + assert res_git.task_id is not None + + # 4. git_branch missing name + task_git_missing = SubAgentTask( + agent_type=SubAgentType.WORKSPACE, + payload={"operation": "git_branch", "action": "create", "path": str(tmp_path)}, + ) + res_git_missing = await agent.run(task_git_missing) + assert res_git_missing.success is False + + +@pytest.mark.asyncio +async def test_spawner_agent_coverage(tmp_path: Path) -> None: + """Test SubAgentSpawner pool initialization, task submission, and status.""" + from collections.abc import Callable + + from agentcli.subagents.base import SubAgent, SubAgentConfig, SubAgentTask, SubAgentType + from agentcli.subagents.bus import MessageBus + from agentcli.subagents.file_ops import FileOpsAgent + from agentcli.subagents.spawner import SubAgentSpawner + + bus = MessageBus() + configs: dict[str, SubAgentConfig] = { + "file_ops": SubAgentConfig(enabled=True, max_concurrent=2), + } + factories: dict[str, Callable[[], SubAgent]] = { + "file_ops": lambda: FileOpsAgent(), + } + + spawner = SubAgentSpawner(config=configs, agent_factories=factories, message_bus=bus) + await spawner.start() + + # 1. Valid task submit + f = tmp_path / "test.txt" + f.write_text("sample content", encoding="utf-8") + task = SubAgentTask( + agent_type=SubAgentType.FILE_OPS, + payload={"operation": "read", "path": str(f)}, + ) + res = await spawner.submit_task(SubAgentType.FILE_OPS, task) + assert res.success is True + + # 2. Status and resource usage + status = await spawner.get_status() + assert "file_ops" in status + resources = spawner.get_resource_usage() + assert "file_ops" in resources + + # 3. Invalid agent type + with pytest.raises(ValueError, match="No pool for agent type"): + await spawner.submit_task(SubAgentType.PLANNER, task) + + await spawner.shutdown() + + + + diff --git a/tests/test_phase35_token_governor.py b/tests/test_phase35_token_governor.py index acf1c11..869d344 100644 --- a/tests/test_phase35_token_governor.py +++ b/tests/test_phase35_token_governor.py @@ -339,3 +339,247 @@ def test_slash_completer_budget_completions() -> None: assert "low" in comps assert "medium" in comps assert "high" in comps + + +def test_adaptive_compressor_edge_cases_and_metrics() -> None: + """Test all edge cases and helper methods in AdaptiveContextCompressor.""" + from agentcli.memory.adaptive_compressor import CompressionMetrics + + # Test CompressionMetrics + metrics = CompressionMetrics() + assert metrics.compression_ratio == 1.0 + assert metrics.tokens_saved == 0 + d = metrics.to_dict() + assert d["runs_count"] == 0 + assert d["compression_ratio"] == 1.0 + + compressor = AdaptiveContextCompressor(keep_recent_turns=2) + + # Empty content tool output + assert compressor.compress_tool_output("") == "" + assert compressor.compress([], max_context_tokens=100) == [] + assert compressor.compress_tier2([]) == [] + t3_empty = compressor.compress_tier3_emergency([]) + assert len(t3_empty) >= 1 + assert "Emergency Context Budget Reset" in (t3_empty[0].content or "") + + # Small history fits budget immediately + short_history = [ + ChatMessage(role="user", content="hello"), + ChatMessage(role="assistant", content="hi"), + ] + res = compressor.compress(short_history, max_context_tokens=10000) + assert res == short_history + + # compress_tier2 with short history (<= recent_msg_count) without and with system msg + assert compressor.compress_tier2(short_history) == short_history + sys_short = [ChatMessage(role="system", content="sys"), *short_history] + assert compressor.compress_tier2(sys_short) == sys_short + + # compress_tier2 with long history and long message previews (>140 chars) + long_msg = "x" * 200 + older = [ChatMessage(role="user", content=long_msg), ChatMessage(role="assistant", content=long_msg)] + recent = [ + ChatMessage(role="user", content="recent 1"), + ChatMessage(role="assistant", content="recent 2"), + ChatMessage(role="user", content="recent 3"), + ChatMessage(role="assistant", content="recent 4"), + ] + t2_res = compressor.compress_tier2([ChatMessage(role="system", content="sys"), *older, *recent], keep_recent=2) + assert len(t2_res) == 1 + 1 + 4 # sys + summary + 4 recent + assert "..." in (t2_res[1].content or "") + + # compress_tier3_emergency without system message + t3_no_sys = compressor.compress_tier3_emergency( + older, user_goal="solve bug", touched_files=["a.py", "b.py"] + ) + assert len(t3_no_sys) == 2 + assert "solve bug" in (t3_no_sys[0].content or "") + assert "a.py" in (t3_no_sys[0].content or "") + + +def test_governor_serialization_and_resets() -> None: + """Test TokenBudgetGovernor serialization, format summary, and resets.""" + governor = TokenBudgetGovernor(max_cost_usd=5.0, max_tokens=100_000) + governor.record_usage( + model="openai/gpt-4o", + prompt_tokens=5000, + completion_tokens=2000, + cached_tokens=1000, + agent_type="planner", + ) + governor.record_usage( + model="google/gemma-4-31b-it:free", + prompt_tokens=1000, + completion_tokens=500, + agent_type="main", + ) + + summary = governor.format_summary() + assert "Token & Cost Budget" in summary + assert "planner" in summary + assert "gpt-4o" in summary + + # Reset + governor.reset() + assert governor.total_tokens == 0 + assert governor.total_cost_usd == 0.0 + assert len(governor.records) == 0 + + +@pytest.mark.asyncio +async def test_cli_budget_slash_command_execution(monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]) -> None: + """Test all /budget subcommand executions through run_chat.""" + import argparse + import os + from unittest.mock import patch + + from agentcli.cli import run_chat + from agentcli.exit_codes import ExitCode + + inputs = [ + "/budget", + "/budget status", + "/budget low", + "/budget tier high", + "/budget tier invalid_tier", + "/budget set $0.50", + "/budget set invalid", + "/budget set", + "/budget max-tokens 50,000", + "/budget max-tokens invalid", + "/budget max-tokens", + "/budget reset", + "/cost", + "/tokens", + "/exit", + ] + + class MockPrompt: + def __init__(self, *args, **kwargs): + self.lines = list(inputs) + + async def get_input_async(self, prompt="you> "): + if self.lines: + return self.lines.pop(0) + return "/exit" + + monkeypatch.setattr("agentcli.cli.InteractivePrompt", MockPrompt) + + args = argparse.Namespace( + model=None, + file=[], + no_agents_md=True, + show_model=False, + resume=None, + allow_write=False, + plain=True, + no_color=True, + budget=None, + max_cost=None, + ) + config = Config() + + with patch.dict(os.environ, {"OPENROUTER_API_KEY": "sk-dummy"}): + exit_code = await run_chat(args, config) + assert exit_code == ExitCode.SUCCESS + + out, _ = capsys.readouterr() + assert "Current budget tier" in out + assert "Budget tier updated to: low" in out + assert "Budget tier updated to: high" in out + assert "Session budget ceiling set to: $0.5000 USD" in out + assert "Session token budget ceiling set to: 50,000 tokens" in out + assert "Budget and cost counters reset" in out + assert "Token Usage" in out + assert "Token & Cost Budget" in out + + +@pytest.mark.asyncio +async def test_tui_phase35_slash_commands(monkeypatch: pytest.MonkeyPatch) -> None: + """Test TUI handles /budget, /cost, /tokens commands seamlessly.""" + from unittest.mock import AsyncMock, MagicMock + + from agentcli.ui.tui_app import TUIApplication + + session = MagicMock() + session.governor = TokenBudgetGovernor(max_cost_usd=1.0) + session.cumulative_cost_usd = 0.02 + session.get_session_stats = AsyncMock( + return_value={"total_tokens": 1000, "user_tokens": 600, "assistant_tokens": 400} + ) + session.config = Config() + session.registry = None + session.router = None + + tui = TUIApplication(config=session.config, session=session) + mock_event = MagicMock() + + # 1. /cost + await tui._handle_slash_command("/cost", "12:00:00", mock_event) + assert any("Token & Cost Budget" in m[1] for m in tui.state.messages) + + # 2. /tokens + await tui._handle_slash_command("/tokens", "12:00:01", mock_event) + assert any("Token Usage:" in m[1] for m in tui.state.messages) + + # 3. /budget status + await tui._handle_slash_command("/budget status", "12:00:02", mock_event) + + # 4. /budget set + await tui._handle_slash_command("/budget set $2.50", "12:00:03", mock_event) + assert session.governor.max_cost_usd == 2.50 + assert any("ceiling set to: $2.5000" in m[1] for m in tui.state.messages) + + # 5. /budget max-tokens + await tui._handle_slash_command("/budget max-tokens 50000", "12:00:04", mock_event) + assert session.governor.max_tokens == 50000 + + # 6. /budget reset + await tui._handle_slash_command("/budget reset", "12:00:05", mock_event) + assert any("counters reset" in m[1] for m in tui.state.messages) + + # 7. /budget tier + await tui._handle_slash_command("/budget tier high", "12:00:06", mock_event) + assert session.config.routing.budget_tier == "high" + + # 8. /budget invalid + await tui._handle_slash_command("/budget invalid_subcommand", "12:00:07", mock_event) + + +def test_adaptive_compressor_all_passes() -> None: + """Explicitly verify passes 1, 2, 3, and 4 in AdaptiveContextCompressor.compress().""" + compressor = AdaptiveContextCompressor(target_budget_ratio=1.0, keep_recent_turns=2) + + # 1. Pass 1 fit (t1 pruned tool output) + compressor_small = AdaptiveContextCompressor(target_budget_ratio=1.0, max_tool_chars=50, keep_recent_turns=2) + tool_history = [ + ChatMessage(role="user", content="run tool"), + ChatMessage(role="assistant", content="```tool\n" + ("output\n" * 100) + "```"), + ] + res_t1 = compressor_small.compress(tool_history, max_context_tokens=80) + assert len(res_t1) == 2 + assert "omitted by context governor" in (res_t1[1].content or "") + + # 2. Pass 2: compress_tier2 with keep_recent=2 + many_turns = [ + ChatMessage(role="user" if i % 2 == 0 else "assistant", content=f"turn {i}: " + ("context payload " * 20)) + for i in range(10) + ] + t2_direct = compressor.compress_tier2(many_turns, keep_recent=2) + assert any("Milestone Context Summary" in (m.content or "") for m in t2_direct) + assert len(t2_direct) == 5 # 1 summary + 4 recent + + # 3. Pass 3: compress_tier2 with keep_recent=1 + t2_tight_direct = compressor.compress_tier2(many_turns, keep_recent=1) + assert any("Milestone Context Summary" in (m.content or "") for m in t2_tight_direct) + assert len(t2_tight_direct) == 3 # 1 summary + 2 recent + + # 4. Pass 4: Tier 3 emergency reset + t3_direct = compressor.compress_tier3_emergency( + many_turns, user_goal="Emergency recovery", touched_files=["main.py"] + ) + assert any("Emergency Context Budget Reset" in (m.content or "") for m in t3_direct) + + + From 918bf7a491ae9c5ee3cd591f8de51ef7a5bc665c Mon Sep 17 00:00:00 2001 From: De-pitcher Date: Mon, 7 Sep 2026 06:11:18 +0100 Subject: [PATCH 3/4] fix(tui): handle mock governor gracefully and add extended branch coverage --- agentcli/ui/tui_app.py | 22 ++++++++++--- tests/test_coverage_boost.py | 64 ++++++++++++++++++++++++++++++++++++ 2 files changed, 81 insertions(+), 5 deletions(-) diff --git a/agentcli/ui/tui_app.py b/agentcli/ui/tui_app.py index c086c97..88cf03a 100644 --- a/agentcli/ui/tui_app.py +++ b/agentcli/ui/tui_app.py @@ -277,8 +277,14 @@ async def _handle_slash_command(self, text: str, timestamp: str, event: KeyPress if not subcmd or subcmd in ("status", "info", "show"): curr = self.config.routing.budget_tier msg = f"Current budget tier: {curr}" - if self.session and hasattr(self.session, "governor"): - msg += "\n" + self.session.governor.format_summary() + gov = getattr(self.session, "governor", None) + if gov is not None and hasattr(gov, "format_summary"): + try: + summary = gov.format_summary() + if isinstance(summary, str): + msg += "\n" + summary + except Exception: + pass self.add_message("system", msg, timestamp) elif subcmd in ("low", "medium", "high", "tier"): tier_name = val if subcmd == "tier" else subcmd @@ -330,7 +336,7 @@ async def _handle_slash_command(self, text: str, timestamp: str, event: KeyPress self.add_message("system", "Budget and cost counters reset for current session.", timestamp) else: self.add_message( - "system", f"Invalid budget subcommand '{subcmd}'. Choose from: low, medium, high, status, set, max-tokens, reset", timestamp + "system", f"Invalid budget tier '{subcmd}'. Choose from: low, medium, high, status, set, max-tokens, reset", timestamp ) return @@ -406,8 +412,14 @@ async def _handle_slash_command(self, text: str, timestamp: str, event: KeyPress f"({stats['user_tokens']} prompt, {stats['assistant_tokens']} completion) | " f"Estimated Cost: ${cost:.6f} USD" ) - if hasattr(self.session, "governor"): - msg_text += "\n" + self.session.governor.format_summary() + gov = getattr(self.session, "governor", None) + if gov is not None and hasattr(gov, "format_summary"): + try: + summary = gov.format_summary() + if isinstance(summary, str): + msg_text += "\n" + summary + except Exception: + pass self.add_message("system", msg_text, timestamp) return diff --git a/tests/test_coverage_boost.py b/tests/test_coverage_boost.py index 5fe79fc..9d1a4e1 100644 --- a/tests/test_coverage_boost.py +++ b/tests/test_coverage_boost.py @@ -561,6 +561,70 @@ async def test_spawner_agent_coverage(tmp_path: Path) -> None: await spawner.shutdown() +@pytest.mark.asyncio +async def test_tui_extended_coverage() -> None: + """Test extended TUI handlers: skills, mcp, worktrees, presets, and modals.""" + from agentcli.config import Config + from agentcli.skills.manifest import SkillManifest, SkillParameter + from agentcli.ui.tui_app import TUIApplication + + config = Config() + mock_session = MagicMock() + mock_session.forced_model = "test-model" + mock_session.get_session_stats = AsyncMock(return_value={"total_tokens": 100, "user_tokens": 50, "assistant_tokens": 50}) + mock_session.cumulative_cost_usd = 0.001 + mock_session.skill_engine = MagicMock() + mock_session.skill_engine.list_available_skills.return_value = [ + {"name": "test_skill", "version": "1.0.0", "description": "Test skill desc", "source_type": "workspace"} + ] + manifest = SkillManifest( + name="test_skill", + version="1.0.0", + description="Test skill desc", + parameters={"query": SkillParameter(name="query", type="string", description="search query")}, + ) + mock_session.skill_engine.loader.get_skill.return_value = manifest + mock_session.skill_engine.loader.has_skill.return_value = True + mock_session.skill_engine.prepare_skill.return_value = (manifest, "rendered goal") + mock_session.skill_engine.execute_skill = AsyncMock(return_value={"success": True, "output": "skill done"}) + + tui = TUIApplication(config=config, session=mock_session) + mock_event = MagicMock() + + # 1. Skills commands + await tui._handle_slash_command("/skills", "12:00:00", mock_event) + assert any("Available Skills" in m[1] for m in tui.state.messages) + + await tui._handle_slash_command("/skill info test_skill", "12:00:01", mock_event) + assert any("Skill [test_skill]" in m[1] for m in tui.state.messages) + + await tui._handle_slash_command("/skill run test_skill query=hello", "12:00:02", mock_event) + assert any("/skill run test_skill" in m[1] for m in tui.state.messages) + + # 2. Worktree / branch commands + mock_session.worktree_manager = MagicMock() + mock_session.worktree_manager.list_worktrees.return_value = [] + await tui._handle_slash_command("/branch list", "12:00:03", mock_event) + assert any("No active Git worktrees" in m[1] for m in tui.state.messages) + + mock_wt = MagicMock() + mock_wt.branch = "feat/test" + mock_wt.path = "/tmp/wt" + mock_wt.base_ref = "main" + mock_session.worktree_manager.create_worktree.return_value = mock_wt + await tui._handle_slash_command("/branch create feat/test", "12:00:04", mock_event) + assert any("Created worktree" in m[1] for m in tui.state.messages) + + mock_session.worktree_manager.compute_diff.return_value = "--- a/file.py\n+++ b/file.py" + await tui._handle_slash_command("/branch diff feat/test", "12:00:05", mock_event) + assert any("--- Diff for feat/test ---" in m[1] for m in tui.state.messages) + + mock_session.worktree_manager.remove_worktree.return_value = True + await tui._handle_slash_command("/branch discard feat/test", "12:00:06", mock_event) + assert any("Pruned and discarded" in m[1] for m in tui.state.messages) + + + From f9203cd118124c7bef32b3af5bb212a4d03bb1bc Mon Sep 17 00:00:00 2001 From: De-pitcher Date: Mon, 7 Sep 2026 06:13:10 +0100 Subject: [PATCH 4/4] fix(tui): clean up governor type check and eliminate blind exceptions --- agentcli/ui/tui_app.py | 22 ++++++++-------------- 1 file changed, 8 insertions(+), 14 deletions(-) diff --git a/agentcli/ui/tui_app.py b/agentcli/ui/tui_app.py index 88cf03a..b0c9285 100644 --- a/agentcli/ui/tui_app.py +++ b/agentcli/ui/tui_app.py @@ -277,14 +277,11 @@ async def _handle_slash_command(self, text: str, timestamp: str, event: KeyPress if not subcmd or subcmd in ("status", "info", "show"): curr = self.config.routing.budget_tier msg = f"Current budget tier: {curr}" + from ..memory.governor import TokenBudgetGovernor + gov = getattr(self.session, "governor", None) - if gov is not None and hasattr(gov, "format_summary"): - try: - summary = gov.format_summary() - if isinstance(summary, str): - msg += "\n" + summary - except Exception: - pass + if isinstance(gov, TokenBudgetGovernor): + msg += "\n" + gov.format_summary() self.add_message("system", msg, timestamp) elif subcmd in ("low", "medium", "high", "tier"): tier_name = val if subcmd == "tier" else subcmd @@ -412,14 +409,11 @@ async def _handle_slash_command(self, text: str, timestamp: str, event: KeyPress f"({stats['user_tokens']} prompt, {stats['assistant_tokens']} completion) | " f"Estimated Cost: ${cost:.6f} USD" ) + from ..memory.governor import TokenBudgetGovernor + gov = getattr(self.session, "governor", None) - if gov is not None and hasattr(gov, "format_summary"): - try: - summary = gov.format_summary() - if isinstance(summary, str): - msg_text += "\n" + summary - except Exception: - pass + if isinstance(gov, TokenBudgetGovernor): + msg_text += "\n" + gov.format_summary() self.add_message("system", msg_text, timestamp) return