chore: fixing race

Signed-off-by: Danny Kopping <[email protected]>

Author: dannykopping

intercept_anthropic_messages_streaming.go

@@ -414,7 +414,7 @@ newStream:
 			prompt = nil
 		}
 
-		if events.hasInitiated() {
+		if events.isStreaming() {
 			// Check if the stream encountered any errors.
 			if streamErr := stream.Err(); streamErr != nil {
 				if isUnrecoverableError(streamErr) {

intercept_openai_chat_streaming.go

@@ -172,7 +172,7 @@ func (i *OpenAIStreamingChatInterception) ProcessRequest(w http.ResponseWriter,
 			})
 		}
 
-		if events.hasInitiated() {
+		if events.isStreaming() {
 			// Check if the stream encountered any errors.
 			if streamErr := stream.Err(); streamErr != nil {
 				if isUnrecoverableError(streamErr) {

streaming.go

@@ -79,12 +79,14 @@ func (s *eventStream) start(w http.ResponseWriter, r *http.Request) {
 			return
 		case ev, open = <-s.eventsCh: // Once closed, the buffered channel will drain all buffered values before showing as closed.
 			if !open {
+				s.logger.Debug(ctx, "events channel closed")
 				return
 			}
 
 			// Initiate the stream once the first event is received.
 			s.initiateOnce.Do(func() {
 				s.initiated.Store(true)
+				s.logger.Debug(ctx, "stream initiated")
 
 				// Send headers for Server-Sent Event stream.
 				w.Header().Set("Content-Type", "text/event-stream")
@@ -185,8 +187,10 @@ func (s *eventStream) Shutdown(shutdownCtx context.Context) error {
 	return err
 }
 
-func (s *eventStream) hasInitiated() bool {
-	return s.initiated.Load()
+// isStreaming checks if the stream has been initiated, or
+// when events are buffered which - when processed - will initiate the stream.
+func (s *eventStream) isStreaming() bool {
+	return s.initiated.Load() || len(s.eventsCh) > 0
 }
 
 // isConnError checks if an error is related to client disconnection or context cancellation.