oats-coder 1.0.2__py3-none-any.whl

oats/session/task_budget.py

@@ -0,0 +1,159 @@
+"""
+Lightweight task-budget tracking for agent loops.
+
+This helps local-model sessions avoid runaway tool churn by monitoring turn
+count, tool-call volume, and repeated tool patterns.
+"""
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import dataclass, field
+from typing import Any
+
+
+def _env_int(name: str, default: int) -> int:
+    try:
+        return int(os.getenv(name, str(default)))
+    except (TypeError, ValueError):
+        return default
+
+
+@dataclass
+class TaskBudgetSnapshot:
+    iteration: int
+    max_iterations: int
+    tool_calls: int
+    max_tool_calls: int
+    repeated_tool_streak: int
+    pressure: str
+    should_stop: bool
+    guidance: str | None = None
+
+
+@dataclass
+class SessionTaskBudget:
+    max_iterations: int = field(default_factory=lambda: _env_int("CODER_MAX_ITERATIONS", 150))
+    max_tool_calls: int = field(default_factory=lambda: _env_int("CODER_MAX_TOOL_CALLS", 300))
+    repeated_tool_limit: int = 3
+    commit_extension_iterations: int = 8
+    _tool_calls: int = 0
+    _history: list[tuple[str, str]] = field(default_factory=list)
+    _committed: bool = False
+    _commit_iteration: int = 0
+    _commit_tool_calls: int = 0
+
+    def record_tool_call(self, tool_name: str, arguments: dict[str, Any]) -> None:
+        normalized = json.dumps(arguments, sort_keys=True, ensure_ascii=True)
+        self._tool_calls += 1
+        self._history.append((tool_name, normalized))
+        if len(self._history) > 24:
+            self._history.pop(0)
+
+    def commit(self, iteration: int) -> None:
+        """Switch to commit mode: stop the discovery loop but allow the model
+        to continue calling tools needed to finalize work (edits, writes, etc.).
+        Budget limits are extended by commit_extension_iterations as a safety
+        net so a stuck model still terminates eventually."""
+        if self._committed:
+            return
+        self._committed = True
+        self._commit_iteration = iteration
+        self._commit_tool_calls = self._tool_calls
+        self._history.clear()
+
+    def snapshot(self, iteration: int) -> TaskBudgetSnapshot:
+        repeated_streak = self._repeated_streak()
+
+        if self._committed:
+            iters_since = iteration - self._commit_iteration
+            calls_since = self._tool_calls - self._commit_tool_calls
+            hard_stop = (
+                iters_since >= self.commit_extension_iterations
+                or calls_since >= self.commit_extension_iterations * 2
+                or repeated_streak >= self.repeated_tool_limit + 2
+            )
+            return TaskBudgetSnapshot(
+                iteration=iteration,
+                max_iterations=self.max_iterations,
+                tool_calls=self._tool_calls,
+                max_tool_calls=self.max_tool_calls,
+                repeated_tool_streak=repeated_streak,
+                pressure="critical",
+                should_stop=hard_stop,
+                guidance=self._build_commit_guidance(iters_since, repeated_streak),
+            )
+
+        iter_ratio = iteration / max(1, self.max_iterations)
+        tool_ratio = self._tool_calls / max(1, self.max_tool_calls)
+
+        if (
+            iteration >= self.max_iterations
+            or self._tool_calls >= self.max_tool_calls
+            or repeated_streak >= self.repeated_tool_limit + 2
+        ):
+            pressure = "critical"
+            should_stop = True
+        elif (
+            iter_ratio >= 0.85
+            or tool_ratio >= 0.85
+            or repeated_streak >= self.repeated_tool_limit
+        ):
+            pressure = "high"
+            should_stop = False
+        elif (
+            iter_ratio >= 0.65
+            or tool_ratio >= 0.65
+            or repeated_streak >= max(2, self.repeated_tool_limit - 2)
+        ):
+            pressure = "medium"
+            should_stop = False
+        else:
+            pressure = "low"
+            should_stop = False
+
+        guidance = None
+        if pressure in {"medium", "high", "critical"}:
+            guidance = self._build_guidance(pressure, repeated_streak)
+
+        return TaskBudgetSnapshot(
+            iteration=iteration,
+            max_iterations=self.max_iterations,
+            tool_calls=self._tool_calls,
+            max_tool_calls=self.max_tool_calls,
+            repeated_tool_streak=repeated_streak,
+            pressure=pressure,
+            should_stop=should_stop,
+            guidance=guidance,
+        )
+
+    def _repeated_streak(self) -> int:
+        if not self._history:
+            return 0
+        streak = 1
+        last = self._history[-1]
+        for item in reversed(self._history[:-1]):
+            if item == last:
+                streak += 1
+            else:
+                break
+        return streak
+
+    def _build_guidance(self, pressure: str, repeated_streak: int) -> str:
+        line = f"# Task: pressure={pressure}, calls={self._tool_calls}, repeat_streak={repeated_streak}"
+        if pressure in ("high", "critical"):
+            line += "\nFinish current subtask. Avoid repeating identical tool calls."
+        return line
+
+    def _build_commit_guidance(self, iters_since: int, repeated_streak: int) -> str:
+        return (
+            f"# COMMIT MODE (discovery budget exhausted; {iters_since} iters since commit, "
+            f"{self._tool_calls} total tool calls, repeat_streak={repeated_streak})\n"
+            "STOP exploring. Do NOT run more searches, greps, finds, or file reads for "
+            "discovery. Use ONLY the context you already have in this conversation to "
+            "complete the user's original task now.\n"
+            "Tool calls are still allowed, but ONLY to commit work the user asked for "
+            "(edits, writes, running the final command, etc.). If information is missing, "
+            "state what you know, state what's missing, and deliver the best answer you "
+            "can with what you have. Do not start new investigation threads."
+        )

oats/session/task_budget.py.AGENT.python.tools.json

@@ -0,0 +1,7 @@
+{
+    "create_SessionTaskBudget": "create a SessionTaskBudget instance to track tool call limits and iteration budgets for an agent loop",
+    "record_tool_call": "record a tool call with its name and arguments to the SessionTaskBudget history for budget tracking",
+    "snapshot": "get a TaskBudgetSnapshot showing current pressure level, tool call count, and whether the agent should stop",
+    "commit": "commit the SessionTaskBudget to switch from discovery mode to commit mode with extended safety limits",
+    "review_TaskBudgetSnapshot": "review a TaskBudgetSnapshot dataclass to inspect iteration count, pressure level, and guidance for the agent"
+}

oats/session/token_budget.py

@@ -0,0 +1,90 @@
+"""
+Lightweight token-budget tracking for long-running sessions.
+
+This is intentionally heuristic rather than tokenizer-accurate. The goal is to
+help the runtime adapt before local-model context pressure becomes a failure.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from oats.session.message import Message
+
+
+@dataclass
+class BudgetSnapshot:
+    estimated_input_tokens: int
+    context_window: int
+    remaining_tokens: int
+    recommended_max_output_tokens: int
+    pressure: str
+
+
+class SessionTokenBudget:
+    """Estimate context pressure and recommend output budgets."""
+
+    def __init__(
+        self,
+        context_window: int,
+        reserve_output_tokens: int = 4096,
+        minimum_output_tokens: int = 768,
+    ) -> None:
+        self.context_window = max(2048, context_window)
+        self.reserve_output_tokens = max(256, reserve_output_tokens)
+        self.minimum_output_tokens = max(256, minimum_output_tokens)
+
+    def snapshot(
+        self,
+        messages: list[Message],
+        requested_max_tokens: int | None = None,
+    ) -> BudgetSnapshot:
+        estimated_input = self._estimate_tokens(messages)
+        remaining = max(0, self.context_window - estimated_input)
+
+        available_for_output = max(
+            self.minimum_output_tokens,
+            remaining - min(self.reserve_output_tokens, max(0, remaining // 3)),
+        )
+        recommended = available_for_output
+        if requested_max_tokens is not None:
+            recommended = min(recommended, requested_max_tokens)
+        recommended = max(self.minimum_output_tokens, recommended)
+
+        ratio = estimated_input / max(1, self.context_window)
+        if ratio >= 0.92:
+            pressure = "critical"
+        elif ratio >= 0.82:
+            pressure = "high"
+        elif ratio >= 0.65:
+            pressure = "medium"
+        else:
+            pressure = "low"
+
+        return BudgetSnapshot(
+            estimated_input_tokens=estimated_input,
+            context_window=self.context_window,
+            remaining_tokens=remaining,
+            recommended_max_output_tokens=recommended,
+            pressure=pressure,
+        )
+
+    def _estimate_tokens(self, messages: list[Message]) -> int:
+        total_chars = 0
+        for msg in messages:
+            total_chars += len(msg.get_text_content() or "")
+            for tc in msg.get_tool_calls():
+                total_chars += len(str(tc.arguments))
+            for tr in msg.get_tool_results():
+                total_chars += len(tr.output or "")
+                total_chars += len(tr.error or "")
+        return total_chars // 4
+
+
+def format_budget_guidance(snapshot: BudgetSnapshot) -> str:
+    """Create a compact prompt section describing current context pressure."""
+    line = (
+        f"# Budget: {snapshot.estimated_input_tokens}/{snapshot.context_window} tokens used, "
+        f"pressure={snapshot.pressure}, max_output={snapshot.recommended_max_output_tokens}"
+    )
+    if snapshot.pressure in ("high", "critical"):
+        line += "\nFinish current work concisely. Avoid long prose."
+    return line

oats/session/token_budget.py.AGENT.python.tools.json

@@ -0,0 +1,7 @@
+{
+  "create_SessionTokenBudget": "create a SessionTokenBudget instance with a given context window size and output token reserves",
+  "call_snapshot": "call snapshot on a SessionTokenBudget to estimate input tokens and get recommended output budget",
+  "call_estimate_tokens": "call _estimate_tokens on a SessionTokenBudget to get a heuristic token count from a message list",
+  "call_format_budget_guidance": "call format_budget_guidance to produce a compact prompt string describing current context pressure",
+  "review_BudgetSnapshot": "review a BudgetSnapshot dataclass to inspect estimated input tokens, remaining tokens, and pressure level"
+}

oats/session/tool_retention.py

@@ -0,0 +1,80 @@
+"""
+Retention policy for tool results stored in session history.
+
+The goal is to keep enough signal for continuation without letting large tool
+outputs dominate long sessions.
+"""
+from __future__ import annotations
+from dataclasses import dataclass
+from oats.tool.registry import ToolResult
+
+
+@dataclass
+class RetainedToolResult:
+    output: str
+    metadata: dict
+
+
+def retain_tool_result(tool_name: str, result: ToolResult) -> RetainedToolResult:
+    """
+    Compress a tool result for session retention.
+
+    The original `ToolResult` still drives the immediate turn. This helper only
+    determines what should be kept in conversation history for future turns.
+    """
+    output = result.output or ""
+    original_length = len(output)
+    retained = output
+
+    if tool_name == "read":
+        retained = _compress_read_output(output)
+    elif tool_name == "grep":
+        retained = _compress_grep_output(output)
+    elif tool_name == "bash":
+        retained = _compress_bash_output(output, result.error)
+    elif tool_name == "lsp":
+        retained = _compress_lsp_output(output)
+    elif len(output) > 4000:
+        retained = _compress_generic(output, head_lines=40, tail_lines=20)
+
+    metadata = dict(result.metadata or {})
+    metadata["retained_output"] = retained
+    metadata["retention_applied"] = retained != output
+    metadata["original_output_chars"] = original_length
+    metadata["retained_output_chars"] = len(retained)
+    return RetainedToolResult(output=retained, metadata=metadata)
+
+
+def _compress_read_output(output: str) -> str:
+    return _compress_generic(output, head_lines=120, tail_lines=40)
+
+
+def _compress_grep_output(output: str) -> str:
+    return _compress_generic(output, head_lines=80, tail_lines=20)
+
+
+def _compress_bash_output(output: str, error: str | None) -> str:
+    if error:
+        return _compress_generic(output, head_lines=120, tail_lines=80)
+    return _compress_generic(output, head_lines=80, tail_lines=40)
+
+
+def _compress_lsp_output(output: str) -> str:
+    return _compress_generic(output, head_lines=120, tail_lines=20)
+
+
+def _compress_generic(output: str, head_lines: int, tail_lines: int) -> str:
+    if not output:
+        return output
+
+    lines = output.splitlines()
+    total = len(lines)
+    keep = head_lines + tail_lines
+    if total <= keep:
+        return output
+
+    head = lines[:head_lines]
+    tail = lines[-tail_lines:] if tail_lines > 0 else []
+    removed = total - len(head) - len(tail)
+    marker = [f"... [{removed} lines omitted for session retention] ..."]
+    return "\n".join(head + marker + tail)

oats/session/tool_retention.py.AGENT.python.tools.json

@@ -0,0 +1,7 @@
+{
+    "retain_tool_result": "compress a tool result for session retention by keeping head and tail lines and omitting the middle",
+    "compress_read_output": "compress read tool output by retaining the first 120 and last 40 lines for session history",
+    "compress_grep_output": "compress grep tool output by retaining the first 80 and last 20 lines for session history",
+    "compress_bash_output": "compress bash tool output by retaining head and tail lines with extra tail when an error is present",
+    "compress_generic": "compress any tool output by keeping a configurable number of head and tail lines and inserting an omission marker"
+}

oats/session/usage.py

@@ -0,0 +1,139 @@
+"""
+Usage tracking for coder2 - aggregates statistics across all sessions.
+
+Tracks lifetime statistics including:
+- Total sessions created
+- Total prompts (messages) sent
+- Total tokens used (prompt + completion) with input/output breakdown
+- Token breakdown by session
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Optional
+from oats.session.session import list_sessions, SessionInfo
+
+
+@dataclass
+class SessionUsageEntry:
+    """Per-session usage entry for detailed breakdown."""
+
+    session_id: str
+    title: str
+    input_tokens: int = 0
+    output_tokens: int = 0
+    total_tokens: int = 0
+    message_count: int = 0
+    created: Optional[datetime] = None
+
+
+@dataclass
+class UsageStats:
+    """Aggregated usage statistics across all sessions."""
+
+    total_sessions: int = 0
+    total_prompts: int = 0  # Total user messages across all sessions
+    total_tokens: int = 0   # Total tokens across all sessions
+    total_input_tokens: int = 0   # Total input (prompt) tokens
+    total_output_tokens: int = 0  # Total output (completion) tokens
+    earliest_session: Optional[datetime] = None
+    latest_session: Optional[datetime] = None
+    sessions: list[SessionUsageEntry] = field(default_factory=list)
+
+    def to_dict(self) -> dict:
+        """Convert to dictionary for display."""
+        return {
+            "total_sessions": self.total_sessions,
+            "total_prompts": self.total_prompts,
+            "total_tokens": self.total_tokens,
+            "total_input_tokens": self.total_input_tokens,
+            "total_output_tokens": self.total_output_tokens,
+            "earliest_session": self.earliest_session.isoformat() if self.earliest_session else None,
+            "latest_session": self.latest_session.isoformat() if self.latest_session else None,
+        }
+
+
+async def get_usage_stats() -> UsageStats:
+    """Calculate aggregated usage statistics across all sessions.
+
+    Returns:
+        UsageStats with lifetime totals for sessions, prompts, and tokens.
+    """
+    sessions = await list_sessions()
+
+    if not sessions:
+        return UsageStats()
+
+    stats = UsageStats(
+        total_sessions=len(sessions),
+        total_prompts=0,
+        total_tokens=0,
+        total_input_tokens=0,
+        total_output_tokens=0,
+        earliest_session=None,
+        latest_session=None,
+    )
+
+    for session_info in sessions:
+        # Accumulate prompt count (message_count includes all messages)
+        stats.total_prompts += session_info.message_count
+
+        # Accumulate token usage
+        stats.total_tokens += session_info.total_tokens
+        stats.total_input_tokens += getattr(session_info, 'total_input_tokens', 0)
+        stats.total_output_tokens += getattr(session_info, 'total_output_tokens', 0)
+
+        # Track per-session breakdown
+        stats.sessions.append(SessionUsageEntry(
+            session_id=session_info.id,
+            title=session_info.title,
+            input_tokens=getattr(session_info, 'total_input_tokens', 0),
+            output_tokens=getattr(session_info, 'total_output_tokens', 0),
+            total_tokens=session_info.total_tokens,
+            message_count=session_info.message_count,
+            created=session_info.time.created,
+        ))
+
+        # Track earliest and latest sessions
+        created = session_info.time.created
+        if stats.earliest_session is None or created < stats.earliest_session:
+            stats.earliest_session = created
+        if stats.latest_session is None or created > stats.latest_session:
+            stats.latest_session = created
+
+    # Sort sessions by most recent first
+    stats.sessions.sort(key=lambda s: s.created or datetime.min, reverse=True)
+
+    return stats
+
+
+def format_tokens(count: int) -> str:
+    """Format token count with K/M/B suffixes for readability."""
+    if count >= 1_000_000_000:
+        return f"{count / 1_000_000_000:.1f}B"
+    elif count >= 1_000_000:
+        return f"{count / 1_000_000:.1f}M"
+    elif count >= 1_000:
+        return f"{count / 1_000:.1f}K"
+    else:
+        return str(count)
+
+
+def format_usage_summary(stats: UsageStats) -> str:
+    """Format usage statistics as a human-readable summary."""
+    lines = [
+        f"  Total Sessions:  {stats.total_sessions}",
+        f"  Total Prompts:   {stats.total_prompts}",
+        f"  Total Tokens:    {format_tokens(stats.total_tokens)}",
+    ]
+    if stats.total_input_tokens or stats.total_output_tokens:
+        lines.append(f"    Input Tokens:  {format_tokens(stats.total_input_tokens)}")
+        lines.append(f"    Output Tokens: {format_tokens(stats.total_output_tokens)}")
+
+    if stats.earliest_session:
+        lines.append(f"  First Session:   {stats.earliest_session.strftime('%Y-%m-%d %H:%M')}")
+    if stats.latest_session:
+        lines.append(f"  Latest Session:  {stats.latest_session.strftime('%Y-%m-%d %H:%M')}")
+
+    return "\n".join(lines)

oats/session/usage.py.AGENT.python.tools.json

@@ -0,0 +1,7 @@
+{
+    "get_usage_stats": "get aggregated usage statistics across all sessions including total prompts tokens and session counts",
+    "UsageStats_to_dict": "convert a UsageStats object to a dictionary for display with session totals and token breakdowns",
+    "format_tokens": "format a token count integer into a human-readable string with K M or B suffixes",
+    "format_usage_summary": "format a UsageStats object into a human-readable summary string with sessions prompts and tokens",
+    "SessionUsageEntry": "create a per-session usage entry dataclass with session id title tokens and message count"
+}