Update app.py

app.py

@@ -30,6 +30,12 @@ app_state: Dict[str, Any] = {
     "files_since_last_state_upload": 0,
 }
 
+# Global file states (client-managed processing state)
+file_states_cache: Dict[str, Any] = {
+    "next_download_index": 0,
+    "file_states": {}
+}
+
 # --- Hugging Face Utility Functions ---
 
 def get_hf_api() -> HfApi:
@@ -88,15 +94,14 @@ def upload_file_to_huggingface(local_path: str, path_in_repo: str):
 
 def load_state():
     """Loads application state from the local state file."""
-    global app_state
+    global app_state, file_states_cache
     if os.path.exists(STATE_FILE_PATH):
         try:
             with open(STATE_FILE_PATH, "r") as f:
-                # Load the state file content, which contains the file_states and next_download_index
+                # Load the complete state file content
                 state_content = json.load(f)
                 
-                # The app_state itself only tracks counters, not the file_states
-                # We only update the counters here if they exist in the state file
+                # Load internal counters
                 if "total_files_uploaded" in state_content:
                     app_state["total_files_uploaded"] = state_content["total_files_uploaded"]
                 if "current_zip_version" in state_content:
@@ -105,21 +110,32 @@ def load_state():
                     app_state["files_since_last_zip"] = state_content["files_since_last_zip"]
                 if "files_since_last_state_upload" in state_content:
                     app_state["files_since_last_state_upload"] = state_content["files_since_last_state_upload"]
+                
+                # Load file_states (client-managed processing state)
+                if "file_states" in state_content:
+                    file_states_cache["file_states"] = state_content["file_states"]
+                if "next_download_index" in state_content:
+                    file_states_cache["next_download_index"] = state_content["next_download_index"]
                     
-            print(f"State counters loaded successfully from {STATE_FILE_PATH}.")
+            print(f"State loaded successfully: app_state={app_state}, file_states_cache={file_states_cache}")
         except Exception as e:
             print(f"Error loading state file: {e}. Using default state.")
     else:
         print("Local state file not found. Using default state.")
 
 def save_state():
-    """Saves application state (counters) to the local state file."""
-    # This function is now only for saving the internal app_state counters, 
-    # not the file_states, which are managed by the client.
+    """Saves both application state counters and file processing states to the local state file."""
+    global app_state, file_states_cache
     try:
+        # Combine both app_state counters and file_states into one JSON
+        complete_state = {
+            **app_state,  # Include all counters
+            **file_states_cache  # Include file_states and next_download_index
+        }
+        
         with open(STATE_FILE_PATH, "w") as f:
-            json.dump(app_state, f, indent=4)
-        print(f"State counters saved successfully to {STATE_FILE_PATH}.")
+            json.dump(complete_state, f, indent=4)
+        print(f"Complete state saved to {STATE_FILE_PATH}.")
     except Exception as e:
         print(f"Error saving state file: {e}")
 
@@ -228,7 +244,7 @@ app = FastAPI(
 @app.post("/upload/")
 async def upload_file(file: UploadFile = File(...)):
     """Upload a file to the server and trigger periodic tasks."""
-    global app_state
+    global app_state, file_states_cache
     
     # 1. Save the file
     try:
@@ -249,23 +265,38 @@ async def upload_file(file: UploadFile = File(...)):
         
         # --- LOGIC FOR STATE FILE HANDLING ---
         if file.filename == STATE_FILE_NAME:
-            print(f"Intercepted state file upload: {STATE_FILE_NAME}. Saving locally only.")
-            # The client's lock/unlock is now persisted locally on the server.
-            # It will only be uploaded to HF when the STATE_UPLOAD_THRESHOLD is met.
-            # We return success immediately.
-            return {"filename": file.filename, "message": "State file successfully saved locally", "state": app_state}
+            print(f"Intercepted state file upload: {STATE_FILE_NAME}. Merging with cache...")
+            # Read the uploaded state file and merge it with our cache
+            try:
+                with open(file_path, "r") as f:
+                    uploaded_state = json.load(f)
+                    
+                # Merge the uploaded state into our cache
+                if "file_states" in uploaded_state:
+                    file_states_cache["file_states"] = uploaded_state["file_states"]
+                if "next_download_index" in uploaded_state:
+                    file_states_cache["next_download_index"] = uploaded_state["next_download_index"]
+                
+                # Save the merged state back to disk
+                save_state()
+                print(f"✅ State file merged: next_download_index={file_states_cache['next_download_index']}, file_states count={len(file_states_cache['file_states'])}")
+            except Exception as e:
+                print(f"Error processing uploaded state file: {e}")
+            
+            return {"filename": file.filename, "message": "State file successfully merged", "state": file_states_cache}
         # --- END LOGIC FOR STATE FILE HANDLING ---
         
         # 2. Update counters only after successful write of a non-state file (transcription JSON)
         app_state["total_files_uploaded"] += 1
         app_state["files_since_last_zip"] += 1
         app_state["files_since_last_state_upload"] += 1
-        save_state() # Save internal counters locally
+        save_state() # Save both app_state and file_states
+        
+        print(f"✅ File uploaded: {file.filename} | files_since_last_zip={app_state['files_since_last_zip']}")
         
         # 3. Check for state file upload threshold (100 non-state file uploads)
         if app_state["files_since_last_state_upload"] >= STATE_UPLOAD_THRESHOLD:
             print(f"State upload threshold ({STATE_UPLOAD_THRESHOLD}) reached. Uploading state file...")
-            # This uploads the local state file (which includes the latest file_states) to HF
             upload_state_file_to_hf()
             
         # 4. Check for zip upload threshold (100 non-state file uploads)
@@ -275,6 +306,7 @@ async def upload_file(file: UploadFile = File(...)):
             if zip_path:
                 upload_zip_to_hf(zip_path)
                 os.remove(zip_path)
+                cleanup_upload_dir()
                 
         return {"filename": file.filename, "message": "File successfully uploaded", "state": app_state}
     except HTTPException:
@@ -330,45 +362,18 @@ async def list_files():
 async def get_state():
     """
     Get the current application state (including file processing status).
-    This endpoint now ensures it fetches the latest state from the local file.
+    Returns both internal counters and client-managed file_states.
     """
-    # 1. Always read the latest state from the local file (which gets updated by client uploads)
-    if os.path.exists(STATE_FILE_PATH):
-        try:
-            with open(STATE_FILE_PATH, "r") as f:
-                state_content = json.load(f)
-                
-                # Ensure the required keys for the client are present
-                if "file_states" not in state_content:
-                    state_content["file_states"] = {}
-                if "next_download_index" not in state_content:
-                    state_content["next_download_index"] = 0
-                
-                print(f"✅ Read state from file: next_download_index={state_content.get('next_download_index')}, file_states={len(state_content.get('file_states', {}))}")
-                return {"state": state_content}
-                
-        except Exception as e:
-            print(f"Error reading local state file for /state/ endpoint: {e}")
-            raise HTTPException(status_code=500, detail="Error processing state file.")
+    global file_states_cache
     
-    # 2. If file doesn't exist, try to download from HF as fallback
-    download_state_file_from_hf()
-    if os.path.exists(STATE_FILE_PATH):
-        try:
-            with open(STATE_FILE_PATH, "r") as f:
-                state_content = json.load(f)
-                
-                if "file_states" not in state_content:
-                    state_content["file_states"] = {}
-                if "next_download_index" not in state_content:
-                    state_content["next_download_index"] = 0
-                    
-                return {"state": state_content}
-        except Exception as e:
-            print(f"Error reading downloaded state file: {e}")
+    # Return the merged state: app_state counters + file_states cache
+    return_state = {
+        **app_state,
+        **file_states_cache
+    }
     
-    # 3. Return default state if file still not found
-    return {"state": {"next_download_index": 0, "file_states": {}}}
+    print(f"✅ Returning state: next_download_index={file_states_cache.get('next_download_index')}, file_states count={len(file_states_cache.get('file_states', {}))}")
+    return {"state": return_state}
 
 # --- Main execution block for testing/running ---
 if __name__ == "__main__":