ant-design · YumoImer · Dec 23, 2024 · Dec 6, 2024 · Dec 6, 2024 · Dec 9, 2024
diff --git a/components/x-stream/index.ts b/components/x-stream/index.ts
@@ -137,11 +137,15 @@ export interface XStreamOptions<Output> {
   transformStream?: TransformStream<string, Output>;
 }
 
+interface XReadableStream<R = any> extends ReadableStream<R> {
+  [Symbol.asyncIterator](): AsyncGenerator<R, void, unknown>;
+}
+
 /**
  * @description Transform Uint8Array binary stream to {@link SSEOutput} by default
  * @warning The `XStream` only support the `utf-8` encoding. More encoding support maybe in the future.
  */
-async function* XStream<Output = SSEOutput>(options: XStreamOptions<Output>) {
+function XStream<Output = SSEOutput>(options: XStreamOptions<Output>) {
   const { readableStream, transformStream } = options;
 
   if (!(readableStream instanceof ReadableStream)) {
@@ -166,18 +170,23 @@ async function* XStream<Output = SSEOutput>(options: XStreamOptions<Output>) {
         .pipeThrough(splitStream())
         .pipeThrough(splitPart());
 
-  const reader = stream.getReader() as ReadableStreamDefaultReader<Output>;
-
-  while (reader instanceof ReadableStreamDefaultReader) {
-    const { value, done } = await reader.read();
-
-    if (done) break;
-
-    if (!value) continue;
+  /** support async iterator */
+  (stream as XReadableStream<Output>)[Symbol.asyncIterator] = async function* () {
+    const reader = this.getReader();
+    try {
+      while (true) {
+        const { done, value } = await reader.read();
+        if (done) {
+          return;
+        }
+        yield value;
+      }
+    } finally {
+      reader.cancel();
+    }
+  };
-  /** support async iterator */
-  (stream as XReadableStream<Output>)[Symbol.asyncIterator] = async function* () {
-    const reader = this.getReader();
-    try {
-      while (true) {
-        const { done, value } = await reader.read();
-        if (done) {
-          return;
-        }
-        yield value;
-      }
-    } finally {
-      reader.cancel();
-    }
-  };
+  /** support async iterator */
+  (stream as XReadableStream<Output>)[Symbol.asyncIterator] = async function* () {
+    const reader = this.getReader();
+    try {
+      while (true) {
+        const { done, value } = await reader.read();
+        if (done) {
+          return;
+        }
+        yield value;
+      }
+    } finally {
+      reader.releaseLock();
+    }
+  };
-  /** support async iterator */
-  (stream as XReadableStream<Output>)[Symbol.asyncIterator] = async function* () {
-    const reader = this.getReader();
-    try {
-      while (true) {
-        const { done, value } = await reader.read();
-        if (done) {
-          return;
-        }
-        yield value;
-      }
-    } finally {
-      reader.cancel();
-    }
-  };
+  /** support async iterator */
+  (stream as XReadableStream<Output>)[Symbol.asyncIterator] = async function* () {
+    const reader = this.getReader();
+    try {
+      while (true) {
+        const { done, value } = await reader.read();
+        if (done) {
+          return;
+        }
+        yield value;
+      }
+    } finally {
+      reader.releaseLock();
+    }
+  };
 
-    // Transformed data through all transform pipes
-    yield value;
-  }
+  return stream as XReadableStream<Output>;
 }
 
 export default XStream;