Additional session fixes

yourbench_space/app.py

@@ -152,19 +152,23 @@ def run_evaluation_pipeline(oauth_token: gr.OAuthToken | None, org_name, eval_na
 
 def init_session(profile: gr.OAuthProfile | None):
     """Update session on load"""
-    if profile is None: # if we running localy or user is not logged in
-        local_uuid = "local" if is_running_locally() else str(uuid.uuid4())
+    if profile is None:
+        username = "local"
     else:
-        local_uuid = USER_ID_SESSION_MAP.get(profile.username, str(uuid.uuid4()))
-
-    manager = MANAGERS.get(local_uuid)
-    if manager and manager.is_running():
-        logger.info(f"Found existing, running session for {local_uuid}")
-        return gr.State(local_uuid, delete_callback=lambda uid: MANAGERS.remove(uid))
-    
-    if profile:
-        USER_ID_SESSION_MAP[profile.username] = local_uuid
+        username = profile.username
     
+    local_uuid = USER_ID_SESSION_MAP.get(username, str(uuid.uuid4()))
+
+    if manager := MANAGERS.get(local_uuid):
+        if manager.is_running():
+            logger.info(f"Found existing running session for {local_uuid}, restoring")
+            return gr.State(local_uuid, delete_callback=lambda uid: MANAGERS.remove(uid))
+        else:
+            logger.info(f"Found existing stale session for {local_uuid}, starting new")
+            MANAGERS.remove(local_uuid)
+            local_uuid = str(uuid.uuid4())
+
+    USER_ID_SESSION_MAP[username] = local_uuid
     MANAGERS.create(local_uuid)
     logger.info(f"Started session for {local_uuid}")
     return gr.State(local_uuid, delete_callback=lambda uid: MANAGERS.remove(uid))