added compaction every x tokens

agent/context_manager/manager.py

@@ -2,14 +2,25 @@
 Context management for conversation history
 """
 
-from litellm import Message
+import asyncio
+
+from litellm import Message, acompletion
 
 
 class ContextManager:
     """Manages conversation context and message history for the agent"""
 
-    def __init__(self):
+    def __init__(
+        self,
+        max_context: int = 180_000,
+        compact_size: float = 0.1,
+        untouched_messages: int = 5,
+    ):
         self.system_prompt = self._load_system_prompt()
+        self.max_context = max_context
+        self.compact_size = int(max_context * compact_size)
+        self.context_length = len(self.system_prompt) // 4
+        self.untouched_messages = untouched_messages
         self.items: list[Message] = [Message(role="system", content=self.system_prompt)]
 
     def _load_system_prompt(self):
@@ -18,27 +29,58 @@ class ContextManager:
         # TODO: get system prompt from jinja template
         return "You are a helpful assistant."
 
-    def add_message(self, message: Message) -> None:
+    def add_message(self, message: Message, token_count: int = None) -> None:
         """Add a message to the history"""
+        if token_count:
+            self.context_length = token_count
+            print(f"DEBUG : token_count = {self.context_length}")
         self.items.append(message)
 
     def get_messages(self) -> list[Message]:
         """Get all messages for sending to LLM"""
         return self.items
 
-    def compact(self, target_size: int) -> None:
+    async def compact(self, model_name: str) -> None:
         """Remove old messages to keep history under target size"""
-        # Keep system prompt (first message) and remove oldest user/assistant messages
-        if len(self.items) <= target_size:
+        if (self.context_length <= self.max_context) or not self.items:
             return
 
-        # Always keep system prompt
         system_msg = (
             self.items[0] if self.items and self.items[0].role == "system" else None
         )
-        messages_to_keep = self.items[-(target_size - 1) :]
 
+        # Don't summarize a certain number of just-preceding messages
+        recent_messages = self.items[-self.untouched_messages :]
+
+        # Summarize everything in between (skip system prompt, skip preceding n)
+        messages_to_summarize = self.items[1 : -self.untouched_messages]
+
+        # improbable, messages would have to very long
+        if not messages_to_summarize:
+            return
+
+        messages_to_summarize.append(
+            Message(
+                role="user",
+                content="Please provide a concise summary of the conversation above, focusing on key decisions, code changes, problems solved, and important context needed for future turns.",
+            )
+        )
+
+        response = await acompletion(
+            model=model_name,
+            messages=messages_to_summarize,
+            max_completion_tokens=self.compact_size,
+        )
+        summarized_message = Message(
+            role="assistant", content=response.choices[0].message.content
+        )
+
+        # Reconstruct: system + summary + recent 2 messages
         if system_msg:
-            self.items = [system_msg] + messages_to_keep
+            self.items = [system_msg, summarized_message] + recent_messages
         else:
-            self.items = messages_to_keep
+            self.items = [summarized_message] + recent_messages
+
+        self.context_length = (
+            len(self.system_prompt) // 4 + response.usage.completion_tokens
+        )

agent/core/agent_loop.py

@@ -1,11 +1,12 @@
-"""
+"""loop
 Main agent implementation with integrated tool system and MCP support
 """
 
 import asyncio
 import json
 
-from litellm import ChatCompletionMessageToolCall, Message, ModelResponse, acompletion
+from litellm import (ChatCompletionMessageToolCall, Message, ModelResponse,
+                     acompletion)
 from lmnr import observe
 
 from agent.config import Config
@@ -58,17 +59,18 @@ class Handlers:
                     tool_choice="auto",
                 )
 
+                # Extract text response, token usage, and tool calls
                 message = response.choices[0].message
-
-                # Extract content and tool calls
                 content = message.content
+                token_count = response.usage.total_tokens
                 tool_calls: list[ToolCall] = message.get("tool_calls", [])
 
                 # If no tool calls, add assistant message and we're done
                 if not tool_calls:
                     if content:
-                        assistant_msg = Message(role="assistant", content=content)
-                        session.context_manager.add_message(assistant_msg)
+                        assistant_msg = Message(
+                            role="assistant", content=content)
+                        session.context_manager.add_message(assistant_msg, token_count)
                         await session.send_event(
                             Event(
                                 event_type="assistant_message",
@@ -81,9 +83,11 @@ class Handlers:
                 # Add assistant message with tool calls to history
                 # LiteLLM will format this correctly for the provider
                 assistant_msg = Message(
-                    role="assistant", content=content, tool_calls=tool_calls
+                    role="assistant",
+                    content=content,
+                    tool_calls=tool_calls,
                 )
-                session.context_manager.add_message(assistant_msg)
+                session.context_manager.add_message(assistant_msg, token_count)
 
                 if content:
                     await session.send_event(
@@ -139,6 +143,18 @@ class Handlers:
                 )
                 break
 
+        old_length = session.context_manager.context_length
+        await session.context_manager.compact(model_name=session.config.model_name)
+        new_length = session.context_manager.context_length
+
+        if new_length != old_length:
+            await session.send_event(
+                Event(
+                    event_type="compacted",
+                    data={"old_tokens": old_length, "new_tokens": new_length},
+                )
+            )
+
         await session.send_event(
             Event(
                 event_type="turn_complete",
@@ -156,14 +172,14 @@ class Handlers:
     @staticmethod
     async def compact(session: Session) -> None:
         """Handle compact (like compact in codex.rs:1317)"""
-        old_size = len(session.context_manager.items)
-        session.context_manager.compact(target_size=10)
-        new_size = len(session.context_manager.items)
+        old_length = session.context_manager.context_length
+        await session.context_manager.compact(model_name=session.config.model_name)
+        new_length = session.context_manager.context_length
 
         await session.send_event(
             Event(
                 event_type="compacted",
-                data={"removed": old_size - new_size, "remaining": new_size},
+                data={"removed": old_length, "remaining": new_length},
             )
         )
 

agent/core/session.py

@@ -34,7 +34,7 @@ class Session:
         event_queue: asyncio.Queue,
         config: Config | None = None,
     ):
-        self.context_manager = ContextManager()
+        self.context_manager = ContextManager(max_context=4_000, compact_size=0.1, untouched_messages=5)
         self.event_queue = event_queue
         self.session_id = str(uuid.uuid4())
         self.config = config or Config(

agent/main.py

@@ -88,6 +88,10 @@ async def event_listener(
                 break
             elif event.event_type == "processing":
                 print("⏳ Processing...", flush=True)
+            elif event.event_type == "compacted":
+                old_tokens = event.data.get("old_tokens", 0) if event.data else 0
+                new_tokens = event.data.get("new_tokens", 0) if event.data else 0
+                print(f"📦 Compacted context: {old_tokens} → {new_tokens} tokens")
             # Silently ignore other events
 
         except asyncio.CancelledError: