From a20a3fbbf865c221dde74c0f5a0c8c4be15aaace Mon Sep 17 00:00:00 2001
From: regischen <58546692+regischen@users.noreply.github.com>
Date: Mon, 22 Apr 2024 14:35:12 +0000
Subject: [PATCH] feat: provide closeStream (#6659)

Upstream https://github.com/toeverything/blocksuite/pull/6836
---
 .../block-suite-editor/ai/copilot-client.ts           |  4 +++-
 .../blocksuite/block-suite-editor/ai/request.ts       | 11 +++++++++++
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/copilot-client.ts b/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/copilot-client.ts
index 4435c5b9ab29..1dcfc99b386d 100644
--- a/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/copilot-client.ts
+++ b/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/copilot-client.ts
@@ -100,13 +100,15 @@ export class CopilotClient {
   async chatText({
     sessionId,
     messageId,
+    signal,
   }: {
     sessionId: string;
     messageId: string;
+    signal?: AbortSignal;
   }) {
     const url = new URL(`${this.backendUrl}/api/copilot/chat/${sessionId}`);
     url.searchParams.set('messageId', messageId);
-    const response = await fetch(url.toString());
+    const response = await fetch(url.toString(), { signal });
     return response.text();
   }
 
diff --git a/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/request.ts b/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/request.ts
index 6ddf2cc88258..76a76982dab2 100644
--- a/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/request.ts
+++ b/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/request.ts
@@ -28,6 +28,7 @@ export type TextToTextOptions = {
   params?: Record<string, string>;
   timeout?: number;
   stream?: boolean;
+  signal?: AbortSignal;
 };
 
 export function createChatSession({
@@ -102,6 +103,7 @@ export function textToText({
   params,
   sessionId,
   stream,
+  signal,
   timeout = TIMEOUT,
 }: TextToTextOptions) {
   if (stream) {
@@ -120,6 +122,15 @@ export function textToText({
           sessionId: message.sessionId,
           messageId: message.messageId,
         });
+        if (signal) {
+          if (signal.aborted) {
+            eventSource.close();
+            return;
+          }
+          signal.onabort = () => {
+            eventSource.close();
+          };
+        }
         for await (const event of toTextStream(eventSource, { timeout })) {
           if (event.type === 'message') {
             yield event.data;