From 2e369997762889e884b0df0d7b9a70e3bd4ddf08 Mon Sep 17 00:00:00 2001
From: Mattt Zmuda <mattt@replicate.com>
Date: Tue, 17 Sep 2024 14:08:10 -0700
Subject: [PATCH] Refactor final goroutine in streamPrediction

---
 stream.go | 33 +++++++++++++++------------------
 1 file changed, 15 insertions(+), 18 deletions(-)

diff --git a/stream.go b/stream.go
index 6fafc14..9883c23 100644
--- a/stream.go
+++ b/stream.go
@@ -222,30 +222,27 @@ func (r *Client) streamPrediction(ctx context.Context, prediction *Prediction, l
 	}()
 
 	go func() {
+		err := g.Wait()
+
 		defer close(sseChan)
 		defer close(errChan)
 
-		for {
-			select {
-			case <-ctx.Done():
+		if err != nil {
+			if errors.Is(err, io.EOF) {
+				// Attempt to reconnect if the connection was closed before the stream was done
+				r.streamPrediction(ctx, prediction, lastEvent, sseChan, errChan)
 				return
-			case <-done:
-				return
-			default:
-				err := g.Wait()
-				if err != nil {
-					if err == io.EOF {
-						// Attempt to reconnect if the connection was closed before the stream was done
-						r.streamPrediction(ctx, prediction, lastEvent, sseChan, errChan)
-						continue
-					}
+			}
 
-					if errors.Is(err, context.Canceled) {
-						return
-					}
+			if errors.Is(err, context.Canceled) {
+				// Context was canceled, simply return
+				return
+			}
 
-					errChan <- err
-				}
+			select {
+			case errChan <- err:
+			default:
+				// errChan is full or closed
 			}
 		}
 	}()