feat: restore missing Import Model List button and add REST bridge for ComfyUI integration

2026-04-30 21:11:24 +02:00
parent b623b686ba
commit 826f0aa2b5
7 changed files with 775 additions and 419 deletions
@@ -10,6 +10,15 @@ Stop wasting time manually moving files! This tool automates your model manageme
 * **Background Mode:** Keep the downloader running in the system tray.
 * **Vault Security:** Encrypted storage for your API keys (Gemini, Hugging Face).

+## 🔌 ComfyUI Integration
+
+To send workflows directly from ComfyUI to the downloader with a single click:
+
+1. Copy or symlink the folder `comfyui-model-downloader` into your `ComfyUI/custom_nodes/` directory.
+2. Restart ComfyUI.
+3. You will now see a **🚀 Downloader** button in the ComfyUI menu (or under the download icon in the new UI v2).
+4. Make sure the Java application is running to receive the workflow.
+
 ## 🚀 Getting Started

 ### Prerequisites
@@ -122,6 +122,19 @@ public class ConfigService {
        return null;
    }

+    public String getComfyUIPath() { return settings.optString("comfyui_path", ""); }
+    public void setComfyUIPath(String path) { settings.put("comfyui_path", path); save(); }
+
+    public String getApiToken() {
+        String token = settings.optString("api_token", "");
+        if (token.isEmpty() && isUnlocked()) {
+            token = java.util.UUID.randomUUID().toString().replace("-", "");
+            settings.put("api_token", token);
+            save();
+        }
+        return token;
+    }
+
    public boolean isUnlocked() {
        return masterPassword != null;
    }
@@ -29,9 +29,14 @@ public class DefaultDownloadManager implements IDownloadManager {
    private ConfigService configService;

    private void safeUpdateStatus(int index, String status, BiConsumer<Integer, String> statusUpdater) {
-        if (completedIndices.contains(index)) return;
+        // Only mark as completed if it's a final state
+        boolean isFinal = status.contains("✅ Finished") || 
+                         status.contains("Stopped") || 
+                         status.contains("Skipped") || 
+                         status.contains("Already exists") ||
+                         status.startsWith("❌ Error");
        
-        if (status.contains("✅ Finished") || status.contains("Stopped") || status.contains("Skipped") || status.contains("Already exists")) {
+        if (isFinal) {
            completedIndices.add(index);
        }
        statusUpdater.accept(index, status);
@@ -43,39 +48,32 @@ public class DefaultDownloadManager implements IDownloadManager {
        isPaused = false;
        this.currentSelection = selectedIndices;
        completedIndices.clear();
+        
+        java.util.concurrent.atomic.AtomicBoolean finishedCalled = new java.util.concurrent.atomic.AtomicBoolean(false);
+        
        new Thread(() -> {
            try {
-                java.util.concurrent.CompletableFuture<?>[] futures = new java.util.concurrent.CompletableFuture[models.size()];
-                for (int i = 0; i < models.size(); i++) {
-                    if (isStopped) {
-                        for (int j = i; j < models.size(); j++) {
-                            futures[j] = java.util.concurrent.CompletableFuture.completedFuture(null);
-                        }
-                        break;
-                    }
-
-                    while (isPaused && !isStopped) {
-                        try {
-                            Thread.sleep(200);
-                        } catch (InterruptedException e) {
-                            Thread.currentThread().interrupt();
-                            break;
-                        }
-                    }
-
+                int size = models.size();
+                java.util.concurrent.CompletableFuture<?>[] futures = new java.util.concurrent.CompletableFuture[size];
+                
+                for (int i = 0; i < size; i++) {
                    final int index = i;
                    
-                    if (!isSelected(index)) { 
-                        safeUpdateStatus(index, "Skipped (Not Selected)", statusUpdater);
-                        futures[i] = java.util.concurrent.CompletableFuture.completedFuture(null);
+                    if (isStopped) {
+                        futures[index] = java.util.concurrent.CompletableFuture.completedFuture(null);
                        continue;
                    }

-                    ModelInfo info = models.get(i);
-                    futures[i] = java.util.concurrent.CompletableFuture.runAsync(() -> {
-                        if (isStopped) { safeUpdateStatus(index, "Stopped", statusUpdater); return; }
-                        if (!isSelected(index)) { safeUpdateStatus(index, "Skipped (Not Selected)", statusUpdater); return; }
-                        if (info.getUrl() == null || info.getUrl().startsWith("MISSING")) { safeUpdateStatus(index, "Skipped (No URL)", statusUpdater); return; }
+                    if (!isSelected(index)) { 
+                        safeUpdateStatus(index, "Skipped (Not Selected)", statusUpdater);
+                        futures[index] = java.util.concurrent.CompletableFuture.completedFuture(null);
+                        continue;
+                    }
+
+                    ModelInfo info = models.get(index);
+                    futures[index] = java.util.concurrent.CompletableFuture.runAsync(() -> {
+                        if (isStopped || !isSelected(index)) return;
+                        
                        try {
                            String subPath = info.getSave_path();
                            if (subPath == null || subPath.isEmpty() || "default".equalsIgnoreCase(subPath)) {  
@@ -87,17 +85,24 @@ public class DefaultDownloadManager implements IDownloadManager {
                            Files.createDirectories(targetDir);
                            downloadWithResume(info, targetDir.resolve(info.getName()), index, statusUpdater);  
                        } catch (Exception e) { 
-                            if (isStopped || !isSelected(index)) {
-                                safeUpdateStatus(index, !isSelected(index) ? "Skipped (Unchecked)" : "Stopped", statusUpdater);
-                            } else {
+                            if (!isStopped && isSelected(index)) {
                                String msg = e.getMessage() != null ? e.getMessage() : e.getClass().getSimpleName();
-                                safeUpdateStatus(index, "Error: " + msg, statusUpdater);
+                                safeUpdateStatus(index, "❌ Error: " + msg, statusUpdater);
                            }
                        }
                    }, executor);
                }
-                java.util.concurrent.CompletableFuture.allOf(futures).join();
-            } finally { onFinished.run(); }
+                
+                // Wait for all tasks to complete (including skipped ones)
+                try {
+                    java.util.concurrent.CompletableFuture.allOf(futures).join();
+                } catch (Exception ignored) {}
+                
+            } finally {
+                if (finishedCalled.compareAndSet(false, true)) {
+                    onFinished.run();
+                }
+            }
        }).start();
    }

@@ -0,0 +1,125 @@
+package de.tki.comfymodels.service.impl;
+
+import com.sun.net.httpserver.HttpExchange;
+import com.sun.net.httpserver.HttpHandler;
+import com.sun.net.httpserver.HttpServer;
+import org.json.JSONObject;
+import org.springframework.stereotype.Service;
+
+import jakarta.annotation.PostConstruct;
+import jakarta.annotation.PreDestroy;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.net.InetSocketAddress;
+import java.nio.charset.StandardCharsets;
+import java.util.function.Consumer;
+
+@Service
+public class RestBridgeService {
+
+    private HttpServer server;
+    private Consumer<String> workflowConsumer;
+    private int port = 12345;
+    private String expectedApiToken;
+
+    public void setPort(int port) {
+        this.port = port;
+    }
+
+    public void setApiToken(String token) {
+        this.expectedApiToken = token;
+    }
+
+    public void setWorkflowConsumer(Consumer<String> consumer) {
+        this.workflowConsumer = consumer;
+    }
+
+    public void startServer() {
+        if (server != null) return;
+        try {
+            server = HttpServer.create(new InetSocketAddress("127.0.0.1", port), 0);
+            server.createContext("/import", new ImportHandler());
+            server.setExecutor(null);
+            server.start();
+            System.out.println("REST Bridge started on port " + port);
+        } catch (IOException e) {
+            System.err.println("Failed to start REST Bridge: " + e.getMessage());
+        }
+    }
+
+    @PreDestroy
+    public void stopServer() {
+        if (server != null) {
+            server.stop(0);
+            server = null;
+        }
+    }
+
+    private class ImportHandler implements HttpHandler {
+        @Override
+        public void handle(HttpExchange exchange) throws IOException {
+            try {
+                // Add CORS headers to EVERY response
+                exchange.getResponseHeaders().add("Access-Control-Allow-Origin", "*");
+                exchange.getResponseHeaders().add("Access-Control-Allow-Methods", "POST, OPTIONS, GET");
+                exchange.getResponseHeaders().add("Access-Control-Allow-Headers", "Content-Type, Authorization");
+
+                // Handle preflight OPTIONS request
+                if ("OPTIONS".equalsIgnoreCase(exchange.getRequestMethod())) {
+                    exchange.sendResponseHeaders(204, -1);
+                    return;
+                }
+
+                // SECURITY: Validate API Token
+                String authHeader = exchange.getRequestHeaders().getFirst("Authorization");
+                if (expectedApiToken != null && !expectedApiToken.isEmpty()) {
+                    if (authHeader == null || !authHeader.equals("Bearer " + expectedApiToken)) {
+                        System.err.println("REST Bridge: Unauthorized request (Invalid or Missing Token)");
+                        exchange.sendResponseHeaders(401, -1);
+                        return;
+                    }
+                }
+
+                if (!"POST".equalsIgnoreCase(exchange.getRequestMethod())) {
+                    exchange.sendResponseHeaders(405, -1);
+                    return;
+                }
+
+                try (InputStream is = exchange.getRequestBody()) {
+                    String body = new String(is.readAllBytes(), StandardCharsets.UTF_8);
+                    
+                    if (workflowConsumer != null) {
+                        workflowConsumer.accept(body);
+                        
+                        String response = "{\"status\": \"success\", \"message\": \"Workflow received\"}";
+                        byte[] responseBytes = response.getBytes(StandardCharsets.UTF_8);
+                        exchange.getResponseHeaders().set("Content-Type", "application/json");
+                        exchange.sendResponseHeaders(200, responseBytes.length);
+                        try (OutputStream os = exchange.getResponseBody()) {
+                            os.write(responseBytes);
+                        }
+                    } else {
+                        String response = "{\"status\": \"error\", \"message\": \"App not ready\"}";
+                        byte[] responseBytes = response.getBytes(StandardCharsets.UTF_8);
+                        exchange.sendResponseHeaders(503, responseBytes.length);
+                        try (OutputStream os = exchange.getResponseBody()) {
+                            os.write(responseBytes);
+                        }
+                    }
+                }
+            } catch (Exception e) {
+                String response = "{\"status\": \"error\", \"message\": \"" + e.getMessage() + "\"}";
+                byte[] responseBytes = response.getBytes(StandardCharsets.UTF_8);
+                try {
+                    exchange.sendResponseHeaders(500, responseBytes.length);
+                    try (OutputStream os = exchange.getResponseBody()) {
+                        os.write(responseBytes);
+                    }
+                } catch (IOException ignored) {}
+            } finally {
+                exchange.close();
+            }
+        }
+    }
+}
@@ -0,0 +1,88 @@
+package de.tki.comfymodels.service;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+
+import java.io.File;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+class InstallationPathTest {
+
+    @TempDir
+    Path tempDir;
+
+    @Test
+    void testStandardLayout() throws IOException {
+        Path customNodes = tempDir.resolve("custom_nodes");
+        Files.createDirectory(customNodes);
+        
+        File inputPath = tempDir.toFile();
+        File customNodesDir = findCustomNodes(inputPath);
+        
+        assertNotNull(customNodesDir);
+        assertEquals(customNodes.toAbsolutePath().toString(), customNodesDir.getAbsolutePath());
+    }
+
+    @Test
+    void testPortableLayout() throws IOException {
+        Path comfyDir = tempDir.resolve("ComfyUI");
+        Files.createDirectory(comfyDir);
+        Path customNodes = comfyDir.resolve("custom_nodes");
+        Files.createDirectory(customNodes);
+        
+        File inputPath = tempDir.toFile();
+        File customNodesDir = findCustomNodes(inputPath);
+        
+        assertNotNull(customNodesDir);
+        assertEquals(customNodes.toAbsolutePath().toString(), customNodesDir.getAbsolutePath());
+    }
+
+    @Test
+    void testDesktopAppLayout() throws IOException {
+        Path resources = tempDir.resolve("resources");
+        Files.createDirectory(resources);
+        Path comfyDir = resources.resolve("ComfyUI");
+        Files.createDirectory(comfyDir);
+        Path customNodes = comfyDir.resolve("custom_nodes");
+        Files.createDirectory(customNodes);
+        
+        File inputPath = tempDir.toFile();
+        File customNodesDir = findCustomNodes(inputPath);
+        
+        assertNotNull(customNodesDir);
+        assertEquals(customNodes.toAbsolutePath().toString(), customNodesDir.getAbsolutePath());
+    }
+
+    @Test
+    void testDirectSelection() throws IOException {
+        Path customNodes = tempDir.resolve("custom_nodes");
+        Files.createDirectory(customNodes);
+        
+        File inputPath = customNodes.toFile();
+        File customNodesDir = findCustomNodes(inputPath);
+        
+        assertNotNull(customNodesDir);
+        assertEquals(customNodes.toAbsolutePath().toString(), customNodesDir.getAbsolutePath());
+    }
+
+    /**
+     * Logic extracted from Main.java for testing
+     */
+    private File findCustomNodes(File inputPath) {
+        File c1 = new File(inputPath, "custom_nodes");
+        File c2 = new File(inputPath, "ComfyUI" + File.separator + "custom_nodes");
+        File c3 = new File(inputPath, "resources" + File.separator + "ComfyUI" + File.separator + "custom_nodes");
+        File c4 = inputPath;
+
+        if (c1.exists() && c1.isDirectory()) return c1;
+        if (c2.exists() && c2.isDirectory()) return c2;
+        if (c3.exists() && c3.isDirectory()) return c3;
+        if (c4.getName().equalsIgnoreCase("custom_nodes") && c4.isDirectory()) return c4;
+        
+        return null;
+    }
+}
@@ -0,0 +1,53 @@
+package de.tki.comfymodels.service;
+
+import de.tki.comfymodels.service.impl.RestBridgeService;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.net.URI;
+import java.net.http.HttpClient;
+import java.net.http.HttpRequest;
+import java.net.http.HttpResponse;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+class RestBridgeServiceTest {
+
+    private RestBridgeService restBridge;
+    private final int testPort = 12346;
+    private final String testToken = "test-token";
+
+    @BeforeEach
+    void setUp() {
+        restBridge = new RestBridgeService();
+        restBridge.setPort(testPort);
+        restBridge.setApiToken(testToken);
+        restBridge.startServer();
+    }
+
+    @AfterEach
+    void tearDown() {
+        restBridge.stopServer();
+    }
+
+    @Test
+    void testWorkflowImport() throws Exception {
+        AtomicReference<String> receivedWorkflow = new AtomicReference<>();
+        restBridge.setWorkflowConsumer(receivedWorkflow::set);
+
+        String testJson = "{\"nodes\": []}";
+        HttpClient client = HttpClient.newHttpClient();
+        HttpRequest request = HttpRequest.newBuilder()
+                .uri(URI.create("http://127.0.0.1:" + testPort + "/import"))
+                .header("Authorization", "Bearer " + testToken)
+                .POST(HttpRequest.BodyPublishers.ofString(testJson))
+                .build();
+
+        HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
+
+        assertEquals(200, response.statusCode());
+        assertEquals(testJson, receivedWorkflow.get());
+    }
+}