From 7e4abaabff3e130ab1e50c8b16842a6b271c5a5f Mon Sep 17 00:00:00 2001
From: Cohee <18619528+Cohee1207@users.noreply.github.com>
Date: Fri, 19 Jul 2024 00:01:24 +0300
Subject: [PATCH] Interrupt AUTO1111 gens on cancel

---
 src/endpoints/stable-diffusion.js | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/src/endpoints/stable-diffusion.js b/src/endpoints/stable-diffusion.js
index 24c393aa9..9f4a4e741 100644
--- a/src/endpoints/stable-diffusion.js
+++ b/src/endpoints/stable-diffusion.js
@@ -323,6 +323,17 @@ router.post('/generate', jsonParser, async (request, response) => {
         const url = new URL(request.body.url);
         url.pathname = '/sdapi/v1/txt2img';
 
+        const controller = new AbortController();
+        request.socket.removeAllListeners('close');
+        request.socket.on('close', function () {
+            if (!response.writableEnded) {
+                const url = new URL(request.body.url);
+                url.pathname = '/sdapi/v1/interrupt';
+                fetch(url, { method: 'POST', headers: { 'Authorization': getBasicAuthHeader(request.body.auth) } });
+            }
+            controller.abort();
+        });
+
         const result = await fetch(url, {
             method: 'POST',
             body: JSON.stringify(request.body),
@@ -331,6 +342,8 @@ router.post('/generate', jsonParser, async (request, response) => {
                 'Authorization': getBasicAuthHeader(request.body.auth),
             },
             timeout: 0,
+            // @ts-ignore
+            signal: controller.signal,
         });
 
         if (!result.ok) {