update context manager

app.py

@@ -589,14 +589,18 @@ with open("prompts.yaml", 'r') as stream:
 # Override CodeAgent run to automatically clear context periodically
 original_code_agent_run = CodeAgent.run
 
-def context_managed_run(self, query, max_steps=None, persist_conversation=None):
-    """Override to periodically clear conversation context after steps"""
+def context_managed_run(self, query, max_steps=None, persist_conversation=None, stream=False, reset=False, additional_args=None):
+    """Override to periodically clear conversation context after steps."""
     # Initialize step counter if not present
     if not hasattr(self, '_context_step_counter'):
         self._context_step_counter = 0
     
-    # Increment counter
-    self._context_step_counter += 1
+    # Reset counter if reset is requested
+    if reset:
+        self._context_step_counter = 0
+    else:
+        # Increment counter
+        self._context_step_counter += 1
     
     # Every 5 steps, perform aggressive context management
     if self._context_step_counter >= 5:
@@ -620,8 +624,8 @@ def context_managed_run(self, query, max_steps=None, persist_conversation=None):
             
             print(f"Conversation history trimmed from {original_length} to {len(self.conversation)} messages")
     
-    # Call the original run method
-    return original_code_agent_run(self, query, max_steps, persist_conversation)
+    # Call the original run method with all parameters (including stream)
+    return original_code_agent_run(self, query, max_steps, persist_conversation, stream=stream, reset=reset, additional_args=additional_args)
 
 # Apply the monkey patch to CodeAgent class
 CodeAgent.run = context_managed_run