From 2453165e12575a36188cdf3aa61971204601fc76 Mon Sep 17 00:00:00 2001
From: Ethanfel <ethan.fel@ts-pc.fr>
Date: Mon, 5 Jan 2026 16:15:43 +0100
Subject: [PATCH] Update __init__.py

---
 __init__.py | 72 +++++++++++++++++++++++++++++------------------------
 1 file changed, 39 insertions(+), 33 deletions(-)

diff --git a/__init__.py b/__init__.py
index cae3123..d0105fc 100644
--- a/__init__.py
+++ b/__init__.py
@@ -3,21 +3,25 @@ import json
 import aiohttp
 from aiohttp import web
 import server
+import folder_paths # <--- USE COMFTYUI PATH MANAGER
 
-# --- Configuration ---
-# Where to save received workflows on the destination machine
-SAVE_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../user/default/workflows"))
-# If the above path doesn't work for your setup, use the output folder:
-# SAVE_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../output/synced_workflows"))
+# --- Improved Path Configuration ---
+# This attempts to save where ComfyUI looks for workflows by default
+base_path = folder_paths.base_path
+# Try the modern 'user/default/workflows' path first
+SAVE_DIR = os.path.join(base_path, "user", "default", "workflows")
+
+# If that folder structure doesn't exist (older comfy), fallback to 'output'
+if not os.path.exists(os.path.dirname(SAVE_DIR)):
+    SAVE_DIR = folder_paths.get_output_directory()
 
 if not os.path.exists(SAVE_DIR):
     os.makedirs(SAVE_DIR)
 
 class NetworkSync:
-    def __init__(self):
-        pass
+    pass
 
-# 1. RECEIVE Endpoint (Runs on the Destination Machine)
+# 1. RECEIVE Endpoint
 @server.PromptServer.instance.routes.post("/network-sync/receive")
 async def receive_workflow(request):
     try:
@@ -25,54 +29,56 @@ async def receive_workflow(request):
         filename = data.get("name", "synced_workflow")
         workflow = data.get("workflow")
 
+        # Basic security: prevent saving outside the folder
+        filename = os.path.basename(filename) 
         if not filename.endswith(".json"):
             filename += ".json"
 
-        # Sanitize filename to prevent directory traversal
-        filename = os.path.basename(filename)
         save_path = os.path.join(SAVE_DIR, filename)
 
+        # Write the file
         with open(save_path, "w", encoding="utf-8") as f:
             json.dump(workflow, f, indent=2)
 
-        print(f"[NetworkSync] Workflow saved to: {save_path}")
-        return web.json_response({"status": "success", "message": f"Saved {filename}"})
+        print(f"✅ [NetworkSync] Received and updated: {filename}")
+        return web.json_response({"status": "success", "message": f"Saved to {save_path}"})
 
     except Exception as e:
-        print(f"[NetworkSync] Error receiving: {e}")
+        print(f"❌ [NetworkSync] Error: {e}")
         return web.json_response({"status": "error", "message": str(e)}, status=500)
 
-# 2. PUSH Endpoint (Runs on the Source Machine)
+# 2. PUSH Endpoint
 @server.PromptServer.instance.routes.post("/network-sync/push")
 async def push_workflow(request):
     try:
         data = await request.json()
-        target_ip = data.get("target_ip")  # e.g., "192.168.1.5:8188"
-        filename = data.get("name")
-        workflow = data.get("workflow")
+        target_ip = data.get("target_ip")
+        
+        # Auto-correct IP formatting
+        if "http" not in target_ip:
+            target_ip = "http://" + target_ip
+        
+        # Ensure port is present if missing (assume default 8188)
+        if ":" not in target_ip.replace("http://", "").replace("https://", ""):
+             target_ip += ":8188"
 
-        if not target_ip:
-            return web.json_response({"status": "error", "message": "No target IP provided"}, status=400)
+        target_url = f"{target_ip}/network-sync/receive"
+        
+        payload = {
+            "name": data.get("name"), 
+            "workflow": data.get("workflow")
+        }
 
-        # Ensure protocol is present
-        if not target_ip.startswith("http"):
-            target_url = f"http://{target_ip}/network-sync/receive"
-        else:
-            target_url = f"{target_ip}/network-sync/receive"
-
-        payload = {"name": filename, "workflow": workflow}
-
-        # Send from Python (Server A) to Python (Server B)
-        async with aiohttp.ClientSession() as session:
+        # Timeout added so ComfyUI doesn't freeze if target is offline
+        timeout = aiohttp.ClientTimeout(total=5) 
+        async with aiohttp.ClientSession(timeout=timeout) as session:
             async with session.post(target_url, json=payload) as resp:
                 if resp.status == 200:
-                    return web.json_response({"status": "success", "message": "Synced successfully!"})
+                    return web.json_response({"status": "success"})
                 else:
-                    text = await resp.text()
-                    return web.json_response({"status": "error", "message": f"Remote error: {text}"}, status=resp.status)
+                    return web.json_response({"status": "error", "message": f"Target error: {resp.status}"})
 
     except Exception as e:
-        print(f"[NetworkSync] Error pushing: {e}")
         return web.json_response({"status": "error", "message": str(e)}, status=500)
 
 NODE_CLASS_MAPPINGS = {}