From d7ea245646c75f2718f66c586e5646a4c9d04ea1 Mon Sep 17 00:00:00 2001
From: machinepainting <machinepainting@proton.me>
Date: Wed, 23 Apr 2025 05:15:42 -0700
Subject: [PATCH] Update monitor_output.py with restart-safe logic and upload
 wait

---
 dropbox_upload.py |  10 ++--
 monitor_output.py | 113 +++++++++++++++++++++++++++++++---------------
 2 files changed, 81 insertions(+), 42 deletions(-)

diff --git a/dropbox_upload.py b/dropbox_upload.py
index be57d23..2d5492b 100644
--- a/dropbox_upload.py
+++ b/dropbox_upload.py
@@ -6,23 +6,23 @@ import requests
 from dotenv import load_dotenv
 
 def get_token():
-    # ✅ 1. Prioritize direct access_token from general environment variables
+    # 1. Prioritize direct access_token from general environment variables
     access_token = os.getenv("DROPBOX_ACCESS_TOKEN")
     if access_token:
         return access_token
 
-    # ✅ 2. Try refresh_token flow from general environment variables
+    # 2. Try refresh_token flow from general environment variables
     refresh_token = os.getenv("DROPBOX_REFRESH_TOKEN")
     client_id     = os.getenv("DROPBOX_APP_KEY")
     client_secret = os.getenv("DROPBOX_APP_SECRET")
 
-    # ✅ 3. If missing, check for RunPod-style prefixed secrets
+    # 3. If missing, check for RunPod-style prefixed secrets
     refresh_token = refresh_token or os.getenv("RUNPOD_SECRET_DROPBOX_REFRESH_TOKEN")
     access_token  = os.getenv("RUNPOD_SECRET_DROPBOX_ACCESS_TOKEN")
     if access_token:
         return access_token
 
-    # 🔁 4. Final fallback: load from .env for local users
+    # 4. Final fallback: load from .env for local users
     if not all([refresh_token, client_id, client_secret]):
         env_path = os.path.join(os.path.dirname(__file__), ".env")
         if os.path.exists(env_path):
@@ -57,4 +57,4 @@ def upload_to_dropbox(local_file_path, dropbox_folder="/Apps/ComfyUI_Output_File
     with open(local_file_path, "rb") as f:
         dbx.files_upload(f.read(), dropbox_path, mode=dropbox.files.WriteMode.overwrite)
 
-    print(f"📦 Uploaded to Dropbox: {dropbox_path}")
+    print(f"📦📤 Uploaded to Dropbox: {dropbox_path}")
diff --git a/monitor_output.py b/monitor_output.py
index a887cca..b89a001 100644
--- a/monitor_output.py
+++ b/monitor_output.py
@@ -1,48 +1,87 @@
 # ComfyUI_DropSendNode/monitor_output.py
 
-import time
 import os
+import time
+import threading
 from watchdog.observers import Observer
 from watchdog.events import FileSystemEventHandler
 from .dropbox_upload import upload_to_dropbox
 
-uploaded_files = set()
+# Global reference to track and restart the monitor
+watcher_observer = None
+watcher_handler = None
 
-def start_monitoring(folder, dropbox_dest_folder):
-    class Handler(FileSystemEventHandler):
-        def on_created(self, event):
-            if event.is_directory:
-                return
-            file_path = event.src_path
 
-            # Wait for file to finish writing
+def wait_for_complete_write(file_path, timeout=10):
+    last_size = -1
+    stable_count = 0
+    checks = 0
+
+    while checks < timeout * 2:
+        try:
+            current_size = os.path.getsize(file_path)
+            if current_size == last_size:
+                stable_count += 1
+                if stable_count >= 2:
+                    return True
+            else:
+                stable_count = 0
+                last_size = current_size
+        except FileNotFoundError:
+            pass
+
+        time.sleep(0.5)
+        checks += 1
+
+    print(f"⚠️ Timeout: File may still be writing: {file_path}")
+    return False
+
+
+class NewFileHandler(FileSystemEventHandler):
+    def __init__(self, dropbox_dest_folder):
+        super().__init__()
+        self.dropbox_dest_folder = dropbox_dest_folder
+
+    def on_created(self, event):
+        if event.is_directory:
+            return
+
+        file_path = event.src_path
+        print(f"📦✅ Detected new file: {file_path}")
+
+        if not wait_for_complete_write(file_path):
+            print(f"⚠️ Skipping upload: File not stable yet: {file_path}")
+            return
+
+        try:
+            upload_to_dropbox(file_path, self.dropbox_dest_folder)
+        except Exception as e:
+            print(f"📦❌ Upload failed: {e}")
+
+
+def start_monitoring(watch_folder, dropbox_dest_folder):
+    global watcher_observer, watcher_handler
+
+    # Stop and restart watcher if already running
+    if watcher_observer and watcher_observer.is_alive():
+        print("🔁 Restarting monitor to update destination folder.")
+        watcher_observer.stop()
+        watcher_observer.join()
+
+    watcher_handler = NewFileHandler(dropbox_dest_folder)
+    watcher_observer = Observer()
+    watcher_observer.schedule(watcher_handler, watch_folder, recursive=False)
+    watcher_observer.start()
+
+    print(f"🔎🗂️ Now watching: {watch_folder}")
+    print(f"📦📥 Upload target: {dropbox_dest_folder}")
+
+    def keep_alive():
+        try:
             while True:
-                try:
-                    with open(file_path, 'rb') as f:
-                        f.read()
-                    break
-                except Exception:
-                    time.sleep(0.5)
+                time.sleep(1)
+        except KeyboardInterrupt:
+            watcher_observer.stop()
+        watcher_observer.join()
 
-            if file_path not in uploaded_files:
-                try:
-                    upload_to_dropbox(file_path, dropbox_dest_folder)
-                    uploaded_files.add(file_path)
-                    print(f"📦✅ Uploaded: {file_path}")
-                except Exception as e:
-                    print(f"📦❌ Upload failed: {e}")
-
-    # Ensure the folder exists
-    if not os.path.exists(folder):
-        os.makedirs(folder)
-
-    observer = Observer()
-    observer.schedule(Handler(), folder, recursive=False)
-    observer.start()
-
-    try:
-        while True:
-            time.sleep(1)
-    except KeyboardInterrupt:
-        observer.stop()
-    observer.join()
+    threading.Thread(target=keep_alive, daemon=True).start()