Refactor final goroutine in streamPrediction

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
 			}
 		}
 	}()