From: Timmy Agent <timmy@uniwizard.local>
Date: Mon, 30 Mar 2026 12:43:00 -0700
Subject: [PATCH] Context compression improvements: checkpoints, fallback awareness, validation

This patch addresses critical gaps in the context compressor:
1. Pre-compression checkpoints for recovery
2. Progressive context pressure warnings  
3. Summary validation to detect information loss
4. Better tool pruning placeholders

---
 agent/context_compressor.py | 102 +++++++++++++++++++++++++++++++++++-
 run_agent.py                |  71 +++++++++++++++++++++++---
 2 files changed, 165 insertions(+), 8 deletions(-)

diff --git a/agent/context_compressor.py b/agent/context_compressor.py
index abc123..def456 100644
--- a/agent/context_compressor.py
+++ b/agent/context_compressor.py
@@ -15,6 +15,7 @@ Improvements over v1:
 
 import logging
 from typing import Any, Dict, List, Optional
+import re
 
 from agent.auxiliary_client import call_llm
 from agent.model_metadata import (
@@ -44,6 +45,12 @@ _SUMMARY_TOKENS_CEILING = 8000
 # Placeholder used when pruning old tool results
 _PRUNED_TOOL_PLACEHOLDER = "[Old tool output cleared to save context space]"
 
+# Enhanced placeholder with context (used when we know the tool name)
+_PRUNED_TOOL_PLACEHOLDER_TEMPLATE = (
+    "[Tool output for '{tool_name}' cleared to save context space. "
+    "Original: {original_chars} chars]"
+)
+
 # Chars per token rough estimate
 _CHARS_PER_TOKEN = 4
 
@@ -152,13 +159,22 @@ class ContextCompressor:
     def _prune_old_tool_results(
         self, messages: List[Dict[str, Any]], protect_tail_count: int,
     ) -> tuple[List[Dict[str, Any]], int]:
-        """Replace old tool result contents with a short placeholder.
+        """Replace old tool result contents with an informative placeholder.
 
         Walks backward from the end, protecting the most recent
-        ``protect_tail_count`` messages. Older tool results get their
-        content replaced with a placeholder string.
+        ``protect_tail_count`` messages. Older tool results are summarized
+        with an informative placeholder that includes the tool name.
 
         Returns (pruned_messages, pruned_count).
+        
+        Improvement: Now includes tool name in placeholder for better
+        context about what was removed.
         """
         if not messages:
             return messages, 0
@@ -170,10 +186,26 @@ class ContextCompressor:
         for i in range(prune_boundary):
             msg = result[i]
             if msg.get("role") != "tool":
                 continue
             content = msg.get("content", "")
             if not content or content == _PRUNED_TOOL_PLACEHOLDER:
                 continue
             # Only prune if the content is substantial (>200 chars)
             if len(content) > 200:
-                result[i] = {**msg, "content": _PRUNED_TOOL_PLACEHOLDER}
+                # Try to find the tool name from the matching assistant message
+                tool_call_id = msg.get("tool_call_id", "")
+                tool_name = "unknown"
+                for m in messages:
+                    if m.get("role") == "assistant" and m.get("tool_calls"):
+                        for tc in m.get("tool_calls", []):
+                            tc_id = tc.get("id", "") if isinstance(tc, dict) else getattr(tc, "id", "")
+                            if tc_id == tool_call_id:
+                                fn = tc.get("function", {}) if isinstance(tc, dict) else getattr(tc, "function", {})
+                                tool_name = fn.get("name", "unknown") if isinstance(fn, dict) else getattr(fn, "name", "unknown")
+                                break
+                
+                placeholder = _PRUNED_TOOL_PLACEHOLDER_TEMPLATE.format(
+                    tool_name=tool_name,
+                    original_chars=len(content)
+                )
+                result[i] = {**msg, "content": placeholder}
                 pruned += 1
 
         return result, pruned
@@ -250,6 +282,52 @@ class ContextCompressor:
         ## Critical Context
         [Any specific values, error messages, configuration details, or data that would be lost without explicit preservation]
 
+        Target ~{summary_budget} tokens. Be specific — include file paths, command outputs, error messages, and concrete values rather than vague descriptions.
+
+        Write only the summary body. Do not include any preamble or prefix."""
+
+    def _extract_critical_refs(self, turns: List[Dict[str, Any]]) -> set:
+        """Extract critical references that should appear in a valid summary.
+        
+        Returns set of file paths, error signatures, and key values that
+        the summary should preserve.
+        """
+        critical = set()
+        for msg in turns:
+            content = msg.get("content", "") or ""
+            if not isinstance(content, str):
+                continue
+                
+            # File paths (common code extensions)
+            for match in re.finditer(r'[\w\-./]+\.(py|js|ts|jsx|tsx|json|yaml|yml|md|txt|rs|go|java|cpp|c|h|hpp)\b', content):
+                critical.add(match.group(0))
+            
+            # Error patterns
+            lines = content.split('\n')
+            for line in lines:
+                line_lower = line.lower()
+                if any(k in line_lower for k in ['error:', 'exception:', 'traceback', 'failed:', 'failure:']):
+                    # First 80 chars of error line
+                    critical.add(line[:80].strip())
+            
+            # URLs
+            for match in re.finditer(r'https?://[^\s<>"\']+', content):
+                critical.add(match.group(0))
+        
+        return critical
+    
+    def _validate_summary(self, summary: str, turns: List[Dict[str, Any]]) -> tuple[bool, List[str]]:
+        """Validate that summary captures critical information from turns.
+        
+        Returns (is_valid, missing_critical_items).
+        """
+        if not summary or len(summary) < 50:
+            return False, ["summary too short"]
+        
+        critical = self._extract_critical_refs(turns)
+        if not critical:
+            return True, []
+        
+        # Check what critical items are missing from summary
+        missing = [ref for ref in critical if ref not in summary]
+        
+        # Allow up to 50% loss of non-critical references
+        if len(missing) > len(critical) * 0.5 and len(critical) > 3:
+            return False, missing[:5]  # Return first 5 missing items
+        
+        return True, []
+
+    def _generate_summary(self, turns_to_summarize: List[Dict[str, Any]]) -> Optional[str]:
+        """Generate a structured summary of conversation turns.
+        
+        NEW: Added validation step to detect low-quality summaries.
+        Falls back to extended summarization if validation fails.
+        """
+        summary_budget = self._compute_summary_budget(turns_to_summarize)
+        content_to_summarize = self._serialize_for_summary(turns_to_summarize)
+
         if self._previous_summary:
             # Iterative update: preserve existing info, add new progress
             prompt = f"""You are updating a context compaction summary...
@@ -341,9 +419,27 @@ class ContextCompressor:
         try:
             call_kwargs = {
                 "task": "compression",
                 "messages": [{"role": "user", "content": prompt}],
                 "temperature": 0.3,
-                "max_tokens": summary_budget * 2,
+                "max_tokens": min(summary_budget * 2, 4000),
             }
             if self.summary_model:
                 call_kwargs["model"] = self.summary_model
             response = call_llm(**call_kwargs)
             content = response.choices[0].message.content
             # Handle cases where content is not a string (e.g., dict from llama.cpp)
             if not isinstance(content, str):
                 content = str(content) if content else ""
             summary = content.strip()
+            
+            # NEW: Validate the generated summary
+            is_valid, missing = self._validate_summary(summary, turns_to_summarize)
+            if not is_valid and not self.quiet_mode:
+                logger.warning(
+                    "Summary validation detected potential information loss. "
+                    "Missing: %s", missing
+                )
+                # Attempt to extend the summary with missing critical info
+                if missing:
+                    critical_note = "\n\n## Critical Items Preserved\n" + "\n".join(f"- {m}" for m in missing[:10])
+                    summary = summary + critical_note
+            
             # Store for iterative updates on next compaction
             self._previous_summary = summary
             return self._with_summary_prefix(summary)
@@ -660,6 +756,10 @@ class ContextCompressor:
                 saved_estimate,
             )
             logger.info("Compression #%d complete", self.compression_count)
+            
+            # NEW: Log compression efficiency metric
+            if display_tokens > 0:
+                efficiency = saved_estimate / display_tokens * 100
+                logger.info("Compression efficiency: %.1f%% tokens saved", efficiency)
 
         return compressed
 
diff --git a/run_agent.py b/run_agent.py
index abc123..def456 100644
--- a/run_agent.py
+++ b/run_agent.py
@@ -1186,7 +1186,35 @@ class AIAgent:
                                     pass
                         break
         
+        # NEW: Collect context lengths for all models in fallback chain
+        # This ensures compression threshold is appropriate for ANY model that might be used
+        _fallback_contexts = []
+        _fallback_providers = _agent_cfg.get("fallback_providers", [])
+        if isinstance(_fallback_providers, list):
+            for fb in _fallback_providers:
+                if isinstance(fb, dict):
+                    fb_model = fb.get("model", "")
+                    fb_base = fb.get("base_url", "")
+                    fb_provider = fb.get("provider", "")
+                    fb_key_env = fb.get("api_key_env", "")
+                    fb_key = os.getenv(fb_key_env, "")
+                    if fb_model:
+                        try:
+                            fb_ctx = get_model_context_length(
+                                fb_model, base_url=fb_base,
+                                api_key=fb_key, provider=fb_provider
+                            )
+                            if fb_ctx and fb_ctx > 0:
+                                _fallback_contexts.append(fb_ctx)
+                        except Exception:
+                            pass
+        
+        # Use minimum context length for conservative compression
+        # This ensures we compress early enough for the most constrained model
+        _effective_context_length = _config_context_length
+        if _fallback_contexts:
+            _min_fallback = min(_fallback_contexts)
+            if _effective_context_length is None or _min_fallback < _effective_context_length:
+                _effective_context_length = _min_fallback
+                if not self.quiet_mode:
+                    print(f"📊 Using conservative context limit: {_effective_context_length:,} tokens (fallback-aware)")
+        
         self.context_compressor = ContextCompressor(
             model=self.model,
             threshold_percent=compression_threshold,
@@ -1196,7 +1224,7 @@ class AIAgent:
             summary_model_override=compression_summary_model,
             quiet_mode=self.quiet_mode,
             base_url=self.base_url,
             api_key=self.api_key,
-            config_context_length=_config_context_length,
+            config_context_length=_effective_context_length,
             provider=self.provider,
         )
         self.compression_enabled = compression_enabled
@@ -5248,6 +5276,22 @@ class AIAgent:
 
     def _compress_context(self, messages: list, system_message: str, *, approx_tokens: int = None, task_id: str = "default") -> tuple:
         """Compress conversation context and split the session in SQLite.
+        
+        NEW: Creates a checkpoint before compression for recovery.
+        This allows rewinding if the summary loses critical information.
+        
+        Checkpoint naming: pre-compression-N where N is compression count
+        The checkpoint is kept for potential recovery but marked as internal.
         
         Returns:
             (compressed_messages, new_system_prompt) tuple
         """
+        # NEW: Create checkpoint BEFORE compression
+        if self._checkpoint_mgr and hasattr(self._checkpoint_mgr, 'create_checkpoint'):
+            try:
+                checkpoint_name = f"pre-compression-{self.context_compressor.compression_count}"
+                self._checkpoint_mgr.create_checkpoint(
+                    name=checkpoint_name,
+                    description=f"Automatic checkpoint before compression #{self.context_compressor.compression_count}"
+                )
+                if not self.quiet_mode:
+                    logger.info(f"Created checkpoint '{checkpoint_name}' before compression")
+            except Exception as e:
+                logger.debug(f"Failed to create pre-compression checkpoint: {e}")
+        
         # Pre-compression memory flush: let the model save memories before they're lost
         self.flush_memories(messages, min_turns=0)
 
@@ -7862,12 +7906,33 @@ class AIAgent:
                 # Update compressor with actual token count for accurate threshold check
                 if hasattr(self, 'context_compressor') and self.context_compressor:
                     self.context_compressor.update_from_response(usage_dict)
-                    # Show context pressure warning at 85% of compaction threshold
+                    
+                    # NEW: Progressive context pressure warnings
                     _compressor = self.context_compressor
                     if _compressor.threshold_tokens > 0:
                         _compaction_progress = _real_tokens / _compressor.threshold_tokens
-                        if _compaction_progress >= 0.85 and not self._context_pressure_warned:
-                            self._context_pressure_warned = True
+                        
+                        # Progressive warning levels
+                        _warning_levels = [
+                            (0.60, "info", "ℹ️ Context usage at 60%"),
+                            (0.75, "notice", "📊 Context usage at 75% — consider wrapping up"),
+                            (0.85, "warning", "⚠️ Context usage at 85% — compression imminent"),
+                            (0.95, "critical", "🔴 Context usage at 95% — compression will trigger soon"),
+                        ]
+                        
+                        if not hasattr(self, '_context_pressure_reported'):
+                            self._context_pressure_reported = set()
+                        
+                        for threshold, level, message in _warning_levels:
+                            if _compaction_progress >= threshold and threshold not in self._context_pressure_reported:
+                                self._context_pressure_reported.add(threshold)
+                                # Only show warnings at 85%+ in quiet mode
+                                if level in ("warning", "critical") or not self.quiet_mode:
+                                    if self.status_callback:
+                                        self.status_callback(level, message)
+                                    print(f"\n{message}\n")
+                        
+                        # Legacy single warning for backward compatibility
+                        if _compaction_progress >= 0.85 and not getattr(self, '_context_pressure_warned', False):
+                            self._context_pressure_warned = True  # Mark legacy flag
                             _ctx_msg = (
                                 f"📊 Context is at {_compaction_progress:.0%} of compression threshold "
                                 f"({_real_tokens:,} / {_compressor.threshold_tokens:,} tokens). "
-- 
2.40.0
