From 2fcb0883ad86e30ca332caadbf965c7017118aea Mon Sep 17 00:00:00 2001
From: Engel Nyst <enyst@users.noreply.github.com>
Date: Thu, 14 Nov 2024 02:20:56 +0100
Subject: [PATCH] fix initialization

---
 openhands/controller/agent_controller.py |  4 +--
 tests/unit/test_truncation.py            | 45 ++++++++++++++++++------
 2 files changed, 36 insertions(+), 13 deletions(-)

diff --git a/openhands/controller/agent_controller.py b/openhands/controller/agent_controller.py
index ad80cdcb385d..2e4c7c599c89 100644
--- a/openhands/controller/agent_controller.py
+++ b/openhands/controller/agent_controller.py
@@ -699,7 +699,7 @@ def _init_history(self):
         events: list[Event] = []
 
         # If we have a truncation point, get first user message and then rest of history
-        if hasattr(self.state, 'truncation_id') and self.state.truncation_id:
+        if hasattr(self.state, 'truncation_id') and self.state.truncation_id > 0:
             # Find first user message from stream
             first_user_msg = next(
                 (
@@ -822,7 +822,7 @@ def _apply_conversation_window(self, events: list[Event]) -> list[Event]:
 
         # Handle first event in truncated history
         if kept_events:
-            i = mid_point
+            i = 0
             while i < len(kept_events):
                 first_event = kept_events[i]
                 if isinstance(first_event, Observation) and first_event.cause:
diff --git a/tests/unit/test_truncation.py b/tests/unit/test_truncation.py
index 14cb1ca0245c..7d03d2f619a5 100644
--- a/tests/unit/test_truncation.py
+++ b/tests/unit/test_truncation.py
@@ -43,12 +43,14 @@ def test_apply_conversation_window_basic(self, mock_event_stream, mock_agent):
         cmd1 = CmdRunAction(command='ls')
         cmd1._id = 2
         obs1 = CmdOutputObservation(command='ls', content='file1.txt', command_id=2)
+        obs1._id = 3
         obs1._cause = 2
 
         cmd2 = CmdRunAction(command='pwd')
-        cmd2._id = 3
-        obs2 = CmdOutputObservation(command='pwd', content='/home', command_id=3)
-        obs2._cause = 3
+        cmd2._id = 4
+        obs2 = CmdOutputObservation(command='pwd', content='/home', command_id=4)
+        obs2._id = 5
+        obs2._cause = 4
 
         events = [first_msg, cmd1, obs1, cmd2, obs2]
 
@@ -83,15 +85,36 @@ def test_context_window_exceeded_handling(self, mock_event_stream, mock_agent):
         first_msg._source = EventSource.USER
         first_msg._id = 1
 
-        cmd1 = CmdRunAction(command='ls')
-        cmd1._id = 2
-        obs1 = CmdOutputObservation(command='ls', content='file1.txt', command_id=2)
-        obs1._id = 3
-        obs1._cause = 2
+        # Add agent question
+        agent_msg = MessageAction(
+            content='What task would you like me to perform?', wait_for_response=True
+        )
+        agent_msg._source = EventSource.AGENT
+        agent_msg._id = 2
 
-        # Set up mock event stream to return our events
-        mock_event_stream.get_events.return_value = [first_msg, cmd1, obs1]
-        controller.state.history = [first_msg, cmd1, obs1]
+        # Add user response
+        user_response = MessageAction(
+            content='Please list all files and show me current directory',
+            wait_for_response=False,
+        )
+        user_response._source = EventSource.USER
+        user_response._id = 3
+
+        cmd1 = CmdRunAction(command='ls')
+        cmd1._id = 4
+        obs1 = CmdOutputObservation(command='ls', content='file1.txt', command_id=4)
+        obs1._id = 5
+        obs1._cause = 4
+
+        # Update mock event stream to include new messages
+        mock_event_stream.get_events.return_value = [
+            first_msg,
+            agent_msg,
+            user_response,
+            cmd1,
+            obs1,
+        ]
+        controller.state.history = [first_msg, agent_msg, user_response, cmd1, obs1]
         original_history_len = len(controller.state.history)
 
         # Simulate ContextWindowExceededError and truncation