test: harden streaming and MCP recovery contracts

Author: loong10k

src/main/java/io/github/easy4j/comfy/cli/ComfyCli.java

@@ -394,9 +394,16 @@ public static class RunOptions {
         private boolean json;
         private boolean jsonStream;
 
+        public RunOptions() { }
+
         public RunOptions(String workflowPath) {
             this.workflowPath = requireNonBlank("workflowPath", workflowPath);
         }
+
+        public RunOptions workflow(String value) {
+            this.workflowPath = value == null ? null : requireNonBlank("workflowPath", value);
+            return this;
+        }
         public RunOptions wait(boolean value) { this.wait = Boolean.valueOf(value); return this; }
         public RunOptions prompt(String value) { this.prompt = value; return this; }
         public RunOptions set(String value) { this.setOverrides.add(requireNonBlank("set", value)); return this; }
@@ -414,6 +421,12 @@ public RunOptions(String workflowPath) {
         public RunOptions jsonStream(boolean value) { this.jsonStream = value; return this; }
 
         List<String> toArgs(String defaultWhere) {
+            if (workflowPath != null && prompt != null) {
+                throw new IllegalStateException("run workflow and prompt are mutually exclusive");
+            }
+            if (workflowPath == null && (prompt == null || prompt.trim().isEmpty())) {
+                throw new IllegalStateException("run requires workflow or prompt");
+            }
             List<String> args = new ArrayList<String>();
             if (json) args.add("--json");
             if (jsonStream) args.add("--json-stream");

src/main/java/io/github/easy4j/comfy/cli/ComfyCliStreamExecutor.java

@@ -128,7 +128,7 @@ public ComfyCliStreamSession execute(final ComfyCliStreamListener listener, Stri
 
     private void readStdout(Process process, ComfyCliStreamListener listener, CappedBytes retained,
                             AtomicReference<Throwable> listenerFailure) {
-        readLines(process.getInputStream(), new LineConsumer() {
+        readLines(process.getInputStream(), effectiveLineLimit(config.getMaxStdoutBytes()), new LineConsumer() {
             @Override public void accept(String line) {
                 retained.append((line + "\n").getBytes(StandardCharsets.UTF_8));
                 if (line.trim().isEmpty()) return;
@@ -150,7 +150,7 @@ private void readStdout(Process process, ComfyCliStreamListener listener, Capped
 
     private void readStderr(Process process, ComfyCliStreamListener listener, CappedBytes retained,
                             AtomicReference<Throwable> listenerFailure) {
-        readLines(process.getErrorStream(), new LineConsumer() {
+        readLines(process.getErrorStream(), effectiveLineLimit(config.getMaxStderrBytes()), new LineConsumer() {
             @Override public void accept(String line) {
                 retained.append((line + "\n").getBytes(StandardCharsets.UTF_8));
                 try {
@@ -163,7 +163,7 @@ private void readStderr(Process process, ComfyCliStreamListener listener, Capped
         });
     }
 
-    private static void readLines(InputStream input, LineConsumer consumer) {
+    private static void readLines(InputStream input, int maxLineBytes, LineConsumer consumer) {
         ByteArrayOutputStream line = new ByteArrayOutputStream();
         try {
             int b;
@@ -172,7 +172,7 @@ private static void readLines(InputStream input, LineConsumer consumer) {
                     consumer.accept(new String(line.toByteArray(), StandardCharsets.UTF_8));
                     line.reset();
                 } else if (b != '\r') {
-                    line.write(b);
+                    if (line.size() < maxLineBytes) line.write(b);
                 }
             }
             if (line.size() > 0) consumer.accept(new String(line.toByteArray(), StandardCharsets.UTF_8));
@@ -182,6 +182,10 @@ private static void readLines(InputStream input, LineConsumer consumer) {
         }
     }
 
+    private static int effectiveLineLimit(int configured) {
+        return configured > 0 ? configured : 16 * 1024 * 1024;
+    }
+
     private void terminate(Process process) {
         if (!process.isAlive()) return;
         process.destroy();

src/test/java/io/github/easy4j/comfy/cli/ComfyCliParityTest.java

@@ -106,6 +106,15 @@ void shouldValidateRoutingAndDynamicParameterNames() {
                 () -> cli().skillsStatus("workspace"));
     }
 
+    @Test
+    void runOptionsShouldSupportPromptOnlyAndRejectPromptWorkflowConflict() {
+        String promptOnly = cli().run(new ComfyCli.RunOptions().prompt("a cat").noWatch(true))
+                .getStdout();
+        assertTrue(promptOnly.contains("run --prompt a cat --no-watch"));
+        assertThrows(IllegalStateException.class,
+                () -> cli().run(new ComfyCli.RunOptions("wf.json").prompt("cat")));
+    }
+
     @Test
     void copiedGenerateOptionsMustBeIndependent() {
         ComfyCli.GenerateOptions original = new ComfyCli.GenerateOptions().prompt("cat").param("quality", "high");

src/test/java/io/github/easy4j/comfy/mcp/ComfyMcpHardeningTest.java

@@ -92,6 +92,44 @@ void notificationsMustBeSurfaced() {
         }
     }
 
+    @Test
+    void elicitationShouldBeHandledOnlyWhenConfigured() {
+        try (ComfyMcpClient client = new ComfyMcpClient(config())) {
+            client.setElicitationHandler(new ComfyMcpElicitationHandler() {
+                @Override public Object handle(String method, com.fasterxml.jackson.databind.JsonNode params) {
+                    Map<String, Object> result = new LinkedHashMap<String, Object>();
+                    result.put("accepted", Boolean.TRUE);
+                    return result;
+                }
+            });
+            client.connect();
+            ComfyMcpCallResult result = client.callTool("elicitation_test", null);
+            assertFalse(result.isError());
+            assertTrue(result.getText().contains("accepted=true"));
+        }
+    }
+
+    @Test
+    void unexpectedChildExitShouldResetStateAndAllowReconnect() {
+        ComfyMcpClient client = new ComfyMcpClient(config());
+        try {
+            client.connect();
+            assertThrows(ComfyException.class, () -> client.callTool("die", null));
+            long deadline = System.nanoTime() + java.util.concurrent.TimeUnit.SECONDS.toNanos(2);
+            while (client.isConnected() && System.nanoTime() < deadline) {
+                try { Thread.sleep(10L); } catch (InterruptedException e) {
+                    Thread.currentThread().interrupt();
+                    break;
+                }
+            }
+            assertFalse(client.isConnected());
+            assertEquals("0.0.0-test", client.connect());
+            assertTrue(client.isConnected());
+        } finally {
+            client.close();
+        }
+    }
+
     @Test
     void nonTextMcpContentMustRemainStructured() {
         try (ComfyMcpClient client = new ComfyMcpClient(config())) {

src/test/resources/fake-mcp-server.py

@@ -71,6 +71,15 @@ def main():
                 send({"jsonrpc": "2.0", "method": "notifications/progress",
                       "params": {"progress": 0.5}})
                 reply(req_id, {"content": [{"type": "text", "text": "ok"}], "isError": False})
+            elif name == "elicitation_test":
+                send({"jsonrpc": "2.0", "id": 9001, "method": "elicitation/create",
+                      "params": {"message": "confirm"}})
+                response = json.loads(sys.stdin.readline())
+                accepted = bool((response.get("result") or {}).get("accepted"))
+                reply(req_id, {"content": [{"type": "text", "text": "accepted=" + str(accepted).lower()}],
+                               "isError": False})
+            elif name == "die":
+                sys.exit(0)
             elif name == "nope":
                 reply(req_id, {"content": [{"type": "text", "text": "unknown tool: " + name}],
                                "isError": True})