fix: token budget tracker for context overflow (#925)

fix: restore _format_error helper for test compatibility (#916)

agent/profile_isolation.py

@@ -1,262 +0,0 @@
-"""
-Profile Session Isolation — #891
-
-Tags sessions with their originating profile and provides
-filtered access so profiles cannot see each other's data.
-
-Current state: All sessions share one state.db with no profile tag.
-This module adds profile tagging and filtered queries.
-
-Usage:
-    from agent.profile_isolation import tag_session, get_profile_sessions, get_active_profile
-    
-    # Tag a new session with the current profile
-    tag_session(session_id, profile_name)
-    
-    # Get sessions for a specific profile
-    sessions = get_profile_sessions("sprint")
-    
-    # Get current active profile
-    profile = get_active_profile()
-"""
-
-import json
-import os
-import sqlite3
-from pathlib import Path
-from typing import Any, Dict, List, Optional
-from datetime import datetime, timezone
-
-HERMES_HOME = Path(os.getenv("HERMES_HOME", str(Path.home() / ".hermes")))
-SESSIONS_DB = HERMES_HOME / "sessions" / "state.db"
-PROFILE_TAGS_FILE = HERMES_HOME / "profile_session_tags.json"
-
-
-def get_active_profile() -> str:
-    """Get the currently active profile name."""
-    config_path = HERMES_HOME / "config.yaml"
-    if config_path.exists():
-        try:
-            import yaml
-            with open(config_path) as f:
-                cfg = yaml.safe_load(f) or {}
-            return cfg.get("active_profile", "default")
-        except Exception:
-            pass
-    
-    # Check environment
-    return os.getenv("HERMES_PROFILE", "default")
-
-
-def _load_tags() -> Dict[str, str]:
-    """Load session-to-profile mapping."""
-    if not PROFILE_TAGS_FILE.exists():
-        return {}
-    try:
-        with open(PROFILE_TAGS_FILE) as f:
-            return json.load(f)
-    except Exception:
-        return {}
-
-
-def _save_tags(tags: Dict[str, str]):
-    """Save session-to-profile mapping."""
-    PROFILE_TAGS_FILE.parent.mkdir(parents=True, exist_ok=True)
-    with open(PROFILE_TAGS_FILE, "w") as f:
-        json.dump(tags, f, indent=2)
-
-
-def tag_session(session_id: str, profile: Optional[str] = None) -> str:
-    """
-    Tag a session with its originating profile.
-    
-    Returns the profile name used.
-    """
-    if profile is None:
-        profile = get_active_profile()
-    
-    tags = _load_tags()
-    tags[session_id] = profile
-    _save_tags(tags)
-    
-    # Also tag in SQLite if available
-    _tag_session_in_db(session_id, profile)
-    
-    return profile
-
-
-def _tag_session_in_db(session_id: str, profile: str):
-    """Add profile tag to SQLite session store."""
-    if not SESSIONS_DB.exists():
-        return
-    
-    try:
-        conn = sqlite3.connect(str(SESSIONS_DB))
-        cursor = conn.cursor()
-        
-        # Check if sessions table has profile column
-        cursor.execute("PRAGMA table_info(sessions)")
-        columns = [row[1] for row in cursor.fetchall()]
-        
-        if "profile" not in columns:
-            # Add profile column
-            cursor.execute("ALTER TABLE sessions ADD COLUMN profile TEXT DEFAULT 'default'")
-        
-        # Update the session's profile
-        cursor.execute(
-            "UPDATE sessions SET profile = ? WHERE session_id = ?",
-            (profile, session_id)
-        )
-        
-        conn.commit()
-        conn.close()
-    except Exception:
-        pass  # SQLite might not be available or schema differs
-
-
-def get_session_profile(session_id: str) -> Optional[str]:
-    """Get the profile that owns a session."""
-    # Check JSON tags first
-    tags = _load_tags()
-    if session_id in tags:
-        return tags[session_id]
-    
-    # Check SQLite
-    if SESSIONS_DB.exists():
-        try:
-            conn = sqlite3.connect(str(SESSIONS_DB))
-            cursor = conn.cursor()
-            cursor.execute(
-                "SELECT profile FROM sessions WHERE session_id = ?",
-                (session_id,)
-            )
-            row = cursor.fetchone()
-            conn.close()
-            if row:
-                return row[0]
-        except Exception:
-            pass
-    
-    return None
-
-
-def get_profile_sessions(
-    profile: Optional[str] = None,
-    limit: int = 100,
-) -> List[Dict[str, Any]]:
-    """
-    Get sessions belonging to a specific profile.
-    
-    Returns list of session dicts.
-    """
-    if profile is None:
-        profile = get_active_profile()
-    
-    sessions = []
-    
-    # Get from JSON tags
-    tags = _load_tags()
-    tagged_sessions = [sid for sid, p in tags.items() if p == profile]
-    
-    # Get from SQLite with profile filter
-    if SESSIONS_DB.exists():
-        try:
-            conn = sqlite3.connect(str(SESSIONS_DB))
-            conn.row_factory = sqlite3.Row
-            cursor = conn.cursor()
-            
-            # Try profile column first
-            try:
-                cursor.execute(
-                    "SELECT * FROM sessions WHERE profile = ? ORDER BY updated_at DESC LIMIT ?",
-                    (profile, limit)
-                )
-                for row in cursor.fetchall():
-                    sessions.append(dict(row))
-            except Exception:
-                # Fallback: filter by tagged session IDs
-                if tagged_sessions:
-                    placeholders = ",".join("?" * len(tagged_sessions[:limit]))
-                    cursor.execute(
-                        f"SELECT * FROM sessions WHERE session_id IN ({placeholders}) ORDER BY updated_at DESC LIMIT ?",
-                        (*tagged_sessions[:limit], limit)
-                    )
-                    for row in cursor.fetchall():
-                        sessions.append(dict(row))
-            
-            conn.close()
-        except Exception:
-            pass
-    
-    return sessions[:limit]
-
-
-def filter_sessions_by_profile(
-    sessions: List[Dict[str, Any]],
-    profile: Optional[str] = None,
-) -> List[Dict[str, Any]]:
-    """Filter a list of sessions to only include those belonging to a profile."""
-    if profile is None:
-        profile = get_active_profile()
-    
-    tags = _load_tags()
-    filtered = []
-    
-    for session in sessions:
-        sid = session.get("session_id") or session.get("id")
-        if not sid:
-            continue
-        
-        # Check tag
-        session_profile = tags.get(sid)
-        if session_profile is None:
-            # Check SQLite
-            session_profile = get_session_profile(sid)
-        
-        if session_profile == profile or session_profile is None:
-            filtered.append(session)
-    
-    return filtered
-
-
-def get_profile_stats() -> Dict[str, Any]:
-    """Get statistics about profile session distribution."""
-    tags = _load_tags()
-    
-    profile_counts = {}
-    for sid, profile in tags.items():
-        profile_counts[profile] = profile_counts.get(profile, 0) + 1
-    
-    total_tagged = len(tags)
-    profiles = list(profile_counts.keys())
-    
-    return {
-        "total_tagged_sessions": total_tagged,
-        "profiles": profiles,
-        "profile_counts": profile_counts,
-        "active_profile": get_active_profile(),
-    }
-
-
-def audit_untagged_sessions() -> List[str]:
-    """Find sessions without a profile tag."""
-    if not SESSIONS_DB.exists():
-        return []
-    
-    try:
-        conn = sqlite3.connect(str(SESSIONS_DB))
-        cursor = conn.cursor()
-        
-        # Get all session IDs
-        cursor.execute("SELECT session_id FROM sessions")
-        all_sessions = {row[0] for row in cursor.fetchall()}
-        conn.close()
-        
-        # Get tagged sessions
-        tags = _load_tags()
-        tagged = set(tags.keys())
-        
-        # Return untagged
-        return list(all_sessions - tagged)
-    except Exception:
-        return []

agent/token_budget.py

@@ -0,0 +1,165 @@
+"""Token Budget — Poka-yoke guard against context overflow.
+
+Progressive warning system with circuit breakers:
+- 60%: Log warning, suggest summarization
+- 80%: Auto-compress, drop raw tool outputs
+- 90%: Block verbose tools, force wrap-up
+- 95%: Graceful termination with summary
+
+Usage:
+    from agent.token_budget import TokenBudget
+    budget = TokenBudget(max_tokens=128000)
+    budget.record_usage(prompt_tokens=500, completion_tokens=200)
+    status = budget.check()
+    # status.level: ok, warning, compress, block, terminate
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Dict, List, Optional
+
+logger = logging.getLogger(__name__)
+
+
+class BudgetLevel(Enum):
+    """Token budget alert levels."""
+    OK = "ok"                 # < 60%
+    WARNING = "warning"       # 60-80%
+    COMPRESS = "compress"     # 80-90%
+    BLOCK = "block"           # 90-95%
+    TERMINATE = "terminate"   # > 95%
+
+
+@dataclass
+class BudgetStatus:
+    """Current budget status."""
+    level: BudgetLevel
+    used_tokens: int
+    max_tokens: int
+    percentage: float
+    remaining: int
+    message: str
+    actions: List[str] = field(default_factory=list)
+
+
+# Default thresholds
+THRESHOLDS = {
+    BudgetLevel.WARNING: 0.60,
+    BudgetLevel.COMPRESS: 0.80,
+    BudgetLevel.BLOCK: 0.90,
+    BudgetLevel.TERMINATE: 0.95,
+}
+
+
+class TokenBudget:
+    """Track token usage and enforce context limits."""
+
+    def __init__(self, max_tokens: int = 128000,
+                 thresholds: Optional[Dict[BudgetLevel, float]] = None):
+        self._max_tokens = max_tokens
+        self._thresholds = thresholds or THRESHOLDS
+        self._prompt_tokens = 0
+        self._completion_tokens = 0
+        self._tool_output_tokens = 0
+        self._history: List[Dict[str, Any]] = []
+
+    @property
+    def used_tokens(self) -> int:
+        return self._prompt_tokens + self._completion_tokens
+
+    @property
+    def remaining(self) -> int:
+        return max(0, self._max_tokens - self.used_tokens)
+
+    @property
+    def percentage(self) -> float:
+        if self._max_tokens == 0:
+            return 0
+        return self.used_tokens / self._max_tokens
+
+    def record_usage(self, prompt_tokens: int = 0, completion_tokens: int = 0,
+                     tool_output_tokens: int = 0):
+        """Record token usage from an API call."""
+        self._prompt_tokens += prompt_tokens
+        self._completion_tokens += completion_tokens
+        self._tool_output_tokens += tool_output_tokens
+        self._history.append({
+            "time": time.time(),
+            "prompt": prompt_tokens,
+            "completion": completion_tokens,
+            "tool_output": tool_output_tokens,
+            "total_used": self.used_tokens,
+        })
+
+    def check(self) -> BudgetStatus:
+        """Check current budget status and return appropriate actions."""
+        pct = self.percentage
+
+        if pct >= self._thresholds.get(BudgetLevel.TERMINATE, 0.95):
+            level = BudgetLevel.TERMINATE
+            msg = f"Context {pct:.0%} full. Session must terminate with summary."
+            actions = ["generate_summary", "terminate_session"]
+        elif pct >= self._thresholds.get(BudgetLevel.BLOCK, 0.90):
+            level = BudgetLevel.BLOCK
+            msg = f"Context {pct:.0%} full. Blocking verbose tool calls."
+            actions = ["block_verbose_tools", "force_wrap_up", "suggest_summary"]
+        elif pct >= self._thresholds.get(BudgetLevel.COMPRESS, 0.80):
+            level = BudgetLevel.COMPRESS
+            msg = f"Context {pct:.0%} full. Auto-compressing conversation."
+            actions = ["auto_compress", "drop_raw_tool_outputs", "suggest_summary"]
+        elif pct >= self._thresholds.get(BudgetLevel.WARNING, 0.60):
+            level = BudgetLevel.WARNING
+            msg = f"Context {pct:.0%} used. Consider summarizing."
+            actions = ["suggest_summary", "log_warning"]
+        else:
+            level = BudgetLevel.OK
+            msg = f"Context OK: {self.used_tokens}/{self._max_tokens} tokens ({pct:.0%})"
+            actions = []
+
+        return BudgetStatus(
+            level=level,
+            used_tokens=self.used_tokens,
+            max_tokens=self._max_tokens,
+            percentage=round(pct, 3),
+            remaining=self.remaining,
+            message=msg,
+            actions=actions,
+        )
+
+    def should_truncate_tool_output(self, estimated_tokens: int) -> bool:
+        """Check if a tool output should be truncated."""
+        if self.used_tokens + estimated_tokens > self._max_tokens * 0.95:
+            return True
+        return False
+
+    def get_truncation_budget(self) -> int:
+        """Get max tokens available for next tool output."""
+        budget = self.remaining - int(self._max_tokens * 0.05)  # Reserve 5%
+        return max(0, budget)
+
+    def reset(self):
+        """Reset budget for new session."""
+        self._prompt_tokens = 0
+        self._completion_tokens = 0
+        self._tool_output_tokens = 0
+        self._history.clear()
+
+    def get_report(self) -> Dict[str, Any]:
+        """Generate usage report."""
+        status = self.check()
+        return {
+            "status": status.level.value,
+            "used_tokens": self.used_tokens,
+            "max_tokens": self._max_tokens,
+            "remaining": self.remaining,
+            "percentage": status.percentage,
+            "prompt_tokens": self._prompt_tokens,
+            "completion_tokens": self._completion_tokens,
+            "tool_output_tokens": self._tool_output_tokens,
+            "message": status.message,
+            "actions": status.actions,
+        }

tests/test_profile_isolation.py

@@ -1,76 +0,0 @@
-"""Tests for profile session isolation (#891)."""
-
-import sys
-import json
-import tempfile
-from pathlib import Path
-
-sys.path.insert(0, str(Path(__file__).parent.parent))
-
-# Override paths for testing
-import agent.profile_isolation as iso_mod
-_test_dir = Path(tempfile.mkdtemp())
-iso_mod.PROFILE_TAGS_FILE = _test_dir / "tags.json"
-
-
-def test_tag_session():
-    """Session gets tagged with profile."""
-    profile = iso_mod.tag_session("sess-1", "sprint")
-    assert profile == "sprint"
-    assert iso_mod.get_session_profile("sess-1") == "sprint"
-
-
-def test_default_profile():
-    """Sessions tagged with default when no profile specified."""
-    profile = iso_mod.tag_session("sess-2")
-    assert profile is not None
-
-
-def test_get_session_profile():
-    """Can retrieve profile for tagged session."""
-    iso_mod.tag_session("sess-3", "fenrir")
-    assert iso_mod.get_session_profile("sess-3") == "fenrir"
-
-
-def test_untagged_returns_none():
-    """Untagged session returns None."""
-    assert iso_mod.get_session_profile("nonexistent") is None
-
-
-def test_profile_stats():
-    """Stats reflect tagged sessions."""
-    iso_mod.tag_session("s1", "default")
-    iso_mod.tag_session("s2", "sprint")
-    iso_mod.tag_session("s3", "sprint")
-    stats = iso_mod.get_profile_stats()
-    assert stats["total_tagged_sessions"] >= 3
-    assert "sprint" in stats["profile_counts"]
-
-
-def test_filter_sessions():
-    """Filter returns only matching profile sessions."""
-    iso_mod.tag_session("filter-1", "alpha")
-    iso_mod.tag_session("filter-2", "beta")
-    iso_mod.tag_session("filter-3", "alpha")
-    
-    sessions = [
-        {"session_id": "filter-1"},
-        {"session_id": "filter-2"},
-        {"session_id": "filter-3"},
-    ]
-    
-    filtered = iso_mod.filter_sessions_by_profile(sessions, "alpha")
-    ids = [s["session_id"] for s in filtered]
-    assert "filter-1" in ids
-    assert "filter-3" in ids
-    assert "filter-2" not in ids
-
-
-if __name__ == "__main__":
-    tests = [test_tag_session, test_default_profile, test_get_session_profile,
-             test_untagged_returns_none, test_profile_stats, test_filter_sessions]
-    for t in tests:
-        print(f"Running {t.__name__}...")
-        t()
-        print("  PASS")
-    print("\nAll tests passed.")

tools/skill_manager_tool.py

@@ -44,6 +44,34 @@ from typing import Dict, Any, Optional, Tuple
 
 logger = logging.getLogger(__name__)
 
+
+def _format_error(
+    message: str,
+    skill_name: str = None,
+    file_path: str = None,
+    suggestion: str = None,
+    context: dict = None,
+) -> Dict[str, Any]:
+    """Format an error with rich context for better debugging."""
+    parts = [message]
+    if skill_name:
+        parts.append(f"Skill: {skill_name}")
+    if file_path:
+        parts.append(f"File: {file_path}")
+    if suggestion:
+        parts.append(f"Suggestion: {suggestion}")
+    if context:
+        for key, value in context.items():
+            parts.append(f"{key}: {value}")
+    return {
+        "success": False,
+        "error": " | ".join(parts),
+        "skill_name": skill_name,
+        "file_path": file_path,
+        "suggestion": suggestion,
+    }
+
+
 # Import security scanner — agent-created skills get the same scrutiny as
 # community hub installs.
 try: