From 124871da76f1ca1a2bd8da1578e6e67c2ee8cacd Mon Sep 17 00:00:00 2001
From: Mikhail <misha_mst@mail.ru>
Date: Tue, 25 Aug 2026 18:02:59 +0300
Subject: [PATCH] Cancel only the owned ComfyUI prompt

---
 adapters/multimodal/client.go      | 39 +++++++++++++++++++++++++++++
 adapters/multimodal/client_test.go | 40 ++++++++++++++++++++++++++++++
 2 files changed, 79 insertions(+)

diff --git a/adapters/multimodal/client.go b/adapters/multimodal/client.go
index d187a74..e35227a 100644
--- a/adapters/multimodal/client.go
+++ b/adapters/multimodal/client.go
@@ -8,6 +8,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"log"
 	"mime/multipart"
 	"net/http"
 	"net/textproto"
@@ -275,6 +276,20 @@ func (c ComfyImageGenerationClient) Generate(ctx context.Context, model, prompt,
 	if accepted.PromptID == "" {
 		return GeneratedImage{}, ErrEmptyResult
 	}
+	// A caller timeout must remove only this prompt. ComfyUI's global
+	// /interrupt endpoint can otherwise stop unrelated users' work and leave
+	// queued requests in an ambiguous state.
+	defer func() {
+		if ctx.Err() != nil {
+			cancelCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+			defer cancel()
+			if err := c.cancelPrompt(cancelCtx, accepted.PromptID); err != nil {
+				// The original context error is more useful to the caller; the
+				// provider cleanup failure remains observable in the server log.
+				log.Printf("comfy prompt cleanup failed for %s: %v", accepted.PromptID, err)
+			}
+		}
+	}()
 
 	interval := c.PollInterval
 	if interval <= 0 {
@@ -306,6 +321,30 @@ func (c ComfyImageGenerationClient) Generate(ctx context.Context, model, prompt,
 	}
 }
 
+func (c ComfyImageGenerationClient) cancelPrompt(ctx context.Context, promptID string) error {
+	if strings.TrimSpace(promptID) == "" {
+		return errors.New("comfy prompt id is empty")
+	}
+	payload, err := json.Marshal(map[string][]string{"delete": {promptID}})
+	if err != nil {
+		return err
+	}
+	request, err := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimRight(c.BaseURL, "/")+"/queue", bytes.NewReader(payload))
+	if err != nil {
+		return err
+	}
+	request.Header.Set("Content-Type", "application/json")
+	response, err := c.client().Do(request)
+	if err != nil {
+		return err
+	}
+	defer response.Body.Close()
+	if response.StatusCode < 200 || response.StatusCode > 299 {
+		return fmt.Errorf("comfy prompt cleanup status %d", response.StatusCode)
+	}
+	return nil
+}
+
 func (c ComfyImageGenerationClient) historyImage(ctx context.Context, promptID string) (GeneratedImage, bool, error) {
 	request, err := http.NewRequestWithContext(ctx, http.MethodGet, strings.TrimRight(c.BaseURL, "/")+"/history/"+url.PathEscape(promptID), nil)
 	if err != nil {
diff --git a/adapters/multimodal/client_test.go b/adapters/multimodal/client_test.go
index 6f1d981..e4ddf49 100644
--- a/adapters/multimodal/client_test.go
+++ b/adapters/multimodal/client_test.go
@@ -3,6 +3,7 @@ package multimodal
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"io"
 	"net/http"
 	"net/http/httptest"
@@ -167,3 +168,42 @@ func TestComfyImageGenerationClientRunsWorkflowAndReturnsImage(t *testing.T) {
 		t.Fatalf("image=%+v err=%v", image, err)
 	}
 }
+
+func TestComfyImageGenerationClientDeletesOnlyCancelledPrompt(t *testing.T) {
+	var deleted []string
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		switch r.URL.Path {
+		case "/prompt":
+			_ = json.NewEncoder(w).Encode(map[string]string{"prompt_id": "cancel-me"})
+		case "/history/cancel-me":
+			<-r.Context().Done()
+		case "/queue":
+			var payload struct {
+				Delete []string `json:"delete"`
+			}
+			if err := json.NewDecoder(r.Body).Decode(&payload); err != nil {
+				t.Fatal(err)
+			}
+			deleted = append(deleted, payload.Delete...)
+			w.WriteHeader(http.StatusNoContent)
+		default:
+			http.NotFound(w, r)
+		}
+	}))
+	defer server.Close()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	done := make(chan error, 1)
+	go func() {
+		_, err := (ComfyImageGenerationClient{BaseURL: server.URL, Client: server.Client(), PollInterval: time.Millisecond, PollTimeout: time.Minute}).Generate(ctx, "", "red bird", "256x256")
+		done <- err
+	}()
+	time.Sleep(10 * time.Millisecond)
+	cancel()
+	if err := <-done; !errors.Is(err, context.Canceled) {
+		t.Fatalf("err=%v", err)
+	}
+	if len(deleted) != 1 || deleted[0] != "cancel-me" {
+		t.Fatalf("deleted=%v", deleted)
+	}
+}