From 34920ea04eb772af899842f772417d35def012e4 Mon Sep 17 00:00:00 2001 From: Engel Nyst Date: Thu, 20 Mar 2025 05:16:49 +0100 Subject: [PATCH] Save agent state (#7372) --- openhands/controller/state/state.py | 25 ++++++++++++++++++------- 1 file changed, 18 insertions(+), 7 deletions(-) diff --git a/openhands/controller/state/state.py b/openhands/controller/state/state.py index 549124fc0e..89b081e72f 100644 --- a/openhands/controller/state/state.py +++ b/openhands/controller/state/state.py @@ -111,12 +111,13 @@ class State: get_conversation_agent_state_filename(sid, user_id), encoded ) - # see if state is in old directory. If yes, delete it. - filename = get_conversation_agent_state_filename(sid) - try: - file_store.delete(filename) - except Exception: - pass + # see if state is in the old directory on saas/remote use cases and delete it. + if user_id: + filename = get_conversation_agent_state_filename(sid) + try: + file_store.delete(filename) + except Exception: + pass except Exception as e: logger.error(f'Failed to save state to session: {e}') raise e @@ -125,6 +126,11 @@ class State: def restore_from_session( sid: str, file_store: FileStore, user_id: str | None = None ) -> 'State': + """ + Restores the state from the previously saved session. + """ + + state: State try: encoded = file_store.read( get_conversation_agent_state_filename(sid, user_id) @@ -132,12 +138,17 @@ class State: pickled = base64.b64decode(encoded) state = pickle.loads(pickled) except FileNotFoundError: + # if user_id is provided, we are in a saas/remote use case + # and we need to check if the state is in the old directory. if user_id: - # see if state is in old directory. If yes, load it. filename = get_conversation_agent_state_filename(sid) encoded = file_store.read(filename) pickled = base64.b64decode(encoded) state = pickle.loads(pickled) + else: + raise FileNotFoundError( + f'Could not restore state from session file for sid: {sid}' + ) except Exception as e: logger.debug(f'Could not restore state from session: {e}') raise e