diff --git a/openhands/controller/agent_controller.py b/openhands/controller/agent_controller.py index 86c663eba2c5..569d37a9be62 100644 --- a/openhands/controller/agent_controller.py +++ b/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. diff --git a/openhands/server/session/agent_session.py b/openhands/server/session/agent_session.py index f198ce8372cd..7d63f9b828d9 100644 --- a/openhands/server/session/agent_session.py +++ b/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 diff --git a/tests/unit/test_truncation.py b/tests/unit/test_truncation.py index 7d03d2f619a5..08e7d8f7be71 100644 --- a/tests/unit/test_truncation.py +++ b/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