Merge pull request #63 from phronmophobic/sse-close-stream

Make sure sse async channel gets closed.

Author: wkok

src/wkok/openai_clojure/sse.clj

@@ -15,7 +15,7 @@
   (when on-next
     (a/go
       (loop []
-        (let [event (a/<! events)]
+        (when-let [event (a/<! events)]
           (when (not= :done event)
             (on-next event)
             (recur)))))))
@@ -43,33 +43,39 @@
 (defn sse-events
   "Returns a core.async channel with events as clojure data structures.
   Inspiration from https://gist.github.com/oliyh/2b9b9107e7e7e12d4a60e79a19d056ee"
-  [{:keys [request params]}]
-  (let [event-stream ^InputStream (:body (http/request (merge request
+  [{:keys [request params] :as m}]
+  (let [close? (:stream/close? params)
+        event-stream ^InputStream (:body (http/request (merge request
                                                               params
                                                               {:as :stream})))
         buffer-size (calc-buffer-size params)
         events (a/chan (a/buffer buffer-size) (map parse-event))]
     (a/thread
-      (loop [byte-coll []]
-        (let [byte-arr (byte-array (max 1 (.available event-stream)))
-              bytes-read (.read event-stream byte-arr)]
-
-          (if (neg? bytes-read)
-
-            ;; Input stream closed, exiting read-loop
-            (.close event-stream)
-
-            (let [next-byte-coll (concat byte-coll (seq byte-arr))
-                  data (slurp (byte-array next-byte-coll))]
-              (if-let [es (not-empty (re-seq event-mask data))]
-                (if (every? true? (map #(a/>!! events %) es))
-                  (recur (drop (apply + (map #(count (.getBytes ^String %)) es))
-                               next-byte-coll))
-
-                  ;; Output stream closed, exiting read-loop
-                  (.close event-stream))
-
-                (recur next-byte-coll)))))))
+      (try
+        (loop [byte-coll []]
+          (let [byte-arr (byte-array (max 1 (.available event-stream)))
+                bytes-read (.read event-stream byte-arr)]
+
+            (if (neg? bytes-read)
+
+              ;; Input stream closed, exiting read-loop
+              nil
+
+              (let [next-byte-coll (concat byte-coll (seq byte-arr))
+                    data (slurp (byte-array next-byte-coll))]
+                (if-let [es (not-empty (re-seq event-mask data))]
+                  (if (every? true? (map #(a/>!! events %) es))
+                    (recur (drop (apply + (map #(count (.getBytes ^String %)) es))
+                                 next-byte-coll))
+
+                    ;; Output stream closed, exiting read-loop
+                    nil)
+
+                  (recur next-byte-coll))))))
+        (finally
+          (when close?
+            (a/close! events))
+          (.close event-stream))))
     events))
 
 (defn sse-request

test/wkok/openai_clojure/sse_test.clj

@@ -11,15 +11,20 @@
 (defn- generate-events
   [data-coll]
   (let [events (map #(str "data: " (json/generate-string %)) data-coll)]
-    (str/join "\n\n" (concat events ["data: [DONE]"]))))
+    (str/join
+     (eduction
+      cat
+      (map #(str % "\n\n"))
+      [events ["data: [DONE]"]]))))
 
 (defn- stream-string
   [^PipedOutputStream output-stream ^String s]
   (future
     (doseq [c (seq (.getBytes s))]
       (.write output-stream ^int c)
       (.flush output-stream)
-      (Thread/sleep 1))))
+      (Thread/sleep 1))
+    (.close output-stream)))
 
 (deftest sse-events-test
   (testing "channel can get events"
@@ -33,8 +38,49 @@
             (is (= (first test-data)
                    (a/<!! events)))
             (is (= (second test-data)
+                   (a/<!! events)))
+            (is (= :done
+                   (a/<!! events)))
+            (is (= nil
+                   (a/poll! events))))))))
+
+  (testing "channel events with `stream/close?` parameter"
+    (let [test-data [{:text "hello"} {:text "world"}]
+          test-events (generate-events test-data)]
+      (with-open [output-stream (PipedOutputStream.)
+                  input-stream (PipedInputStream. output-stream)]
+        (with-redefs [http/request (constantly {:body input-stream})]
+          (let [events (sse/sse-events {:params {:stream/close? true}})]
+            (stream-string output-stream test-events)
+            (is (= (first test-data)
+                   (a/<!! events)))
+            (is (= (second test-data)
+                   (a/<!! events)))
+            (is (= :done
+                   (a/<!! events)))
+            (is (= nil
                    (a/<!! events))))))))
 
+  (testing "channel events with `:on-next`"
+    (doseq [close? [true false]]
+      (let [test-data [{:text "hello"} {:text "world"}]
+            test-events (generate-events test-data)]
+        (with-open [output-stream (PipedOutputStream.)
+                    input-stream (PipedInputStream. output-stream)]
+          (with-redefs [http/request (constantly {:body input-stream})]
+            (let [events (sse/sse-events {:stream/close? close?})
+                  results (promise)
+                  ;; accumulate results and fulfill promise when done
+                  on-next (let [acc (volatile! [])]
+                            (fn [x]
+                              (vswap! acc conj x)
+                              (when (= x (last test-data))
+                                (deliver results @acc))))]
+              (sse/deliver-events events {:on-next on-next})
+              (stream-string output-stream test-events)
+              (is (= test-data
+                     @results))))))))
+
   (testing "support multibytes"
     (let [test-data [{:text "こんにちは"} {:text "你好"}]
           test-events (generate-events test-data)]