Fix history loading when state was corrupt/non-existent (#5946)

Co-authored-by: openhands <[email protected]>

openhands/controller/agent_controller.py

@@ -730,12 +730,20 @@ def set_initial_state(
         # - the previous session, in which case it has history
         # - from a parent agent, in which case it has no history
         # - None / a new state
+
+        # If state is None, we create a brand new state and still load the event stream so we can restore the history
         if state is None:
             self.state = State(
                 inputs={},
                 max_iterations=max_iterations,
                 confirmation_mode=confirmation_mode,
             )
+            self.state.start_id = 0
+
+            self.log(
+                'debug',
+                f'AgentController {self.id} - created new state. start_id: {self.state.start_id}',
+            )
         else:
             self.state = state
 
@@ -747,7 +755,8 @@ def set_initial_state(
                 f'AgentController {self.id} initializing history from event {self.state.start_id}',
             )
 
-            self._init_history()
+        # Always load from the event stream to avoid losing history
+        self._init_history()
 
     def _init_history(self) -> None:
         """Initializes the agent's history from the event stream.

openhands/server/session/agent_session.py

@@ -269,11 +269,26 @@ def _create_controller(
             headless_mode=False,
             status_callback=self._status_callback,
         )
+
+        # Note: We now attempt to restore the state from session here,
+        # but if it fails, we fall back to None and still initialize the controller
+        # with a fresh state. That way, the controller will always load events from the event stream
+        # even if the state file was corrupt.
+
+        restored_state = None
         try:
-            agent_state = State.restore_from_session(self.sid, self.file_store)
-            controller.set_initial_state(agent_state, max_iterations, confirmation_mode)
-            logger.debug(f'Restored agent state from session, sid: {self.sid}')
+            restored_state = State.restore_from_session(self.sid, self.file_store)
         except Exception as e:
-            logger.debug(f'State could not be restored: {e}')
+            if self.event_stream.get_latest_event_id() > 0:
+                # if we have events, we should have a state
+                logger.warning(f'State could not be restored: {e}')
+
+        # Set the initial state through the controller.
+        controller.set_initial_state(restored_state, max_iterations, confirmation_mode)
+        if restored_state:
+            logger.debug(f'Restored agent state from session, sid: {self.sid}')
+        else:
+            logger.debug('New session state created.')
+
         logger.debug('Agent controller initialized.')
         return controller

tests/unit/test_truncation.py

@@ -13,6 +13,8 @@ def mock_event_stream():
     stream = MagicMock()
     # Mock get_events to return an empty list by default
     stream.get_events.return_value = []
+    # Mock get_latest_event_id to return a valid integer
+    stream.get_latest_event_id.return_value = 0
     return stream