first commit

core/providers/utils/stream.go

@@ -0,0 +1,75 @@
+package utils
+
+import (
+	"context"
+
+	schemas "github.com/maximhq/bifrost/core/schemas"
+)
+
+// CheckFirstStreamChunkForError reads the first chunk from a streaming channel to detect
+// errors returned inside HTTP 200 SSE streams (e.g., providers that send rate limit
+// errors as SSE events instead of HTTP 429).
+//
+// If the first chunk is an error, it drains the source channel in the background
+// (so the provider goroutine can exit cleanly) and returns the error for synchronous
+// handling, enabling retries and fallbacks. The returned drainDone channel is closed
+// once the drain completes — callers must wait on it before releasing any resources
+// (e.g., plugin pipelines) that the provider goroutine's postHookRunner may still reference.
+//
+// If the first chunk is valid data, it returns a wrapped channel that re-emits
+// the first chunk followed by all remaining chunks from the source. drainDone is
+// closed when the wrapper goroutine finishes forwarding the source stream.
+//
+// If the source channel is closed immediately (empty stream), it returns a
+// nil channel with nil error. drainDone is already closed.
+//
+// The ctx argument cancels the background forwarding goroutine if the consumer
+// abandons the returned wrapped channel. On ctx.Done the goroutine drains the
+// source stream so the upstream provider's blocked send can exit cleanly.
+func CheckFirstStreamChunkForError(
+	ctx context.Context,
+	stream chan *schemas.BifrostStreamChunk,
+) (chan *schemas.BifrostStreamChunk, <-chan struct{}, *schemas.BifrostError) {
+	firstChunk, ok := <-stream
+	if !ok {
+		// Channel closed immediately (empty stream) — return nil so callers
+		// can distinguish this from a live stream channel.
+		done := make(chan struct{})
+		close(done)
+		return nil, done, nil
+	}
+
+	// Check if first chunk is an error
+	if firstChunk.BifrostError != nil && firstChunk.BifrostError.Error != nil &&
+		(firstChunk.BifrostError.Error.Message != "" || firstChunk.BifrostError.Error.Code != nil || firstChunk.BifrostError.Error.Type != nil) {
+		// Drain source channel to let the provider goroutine exit cleanly
+		done := make(chan struct{})
+		go func() {
+			defer close(done)
+			for range stream {
+			}
+		}()
+		return nil, done, firstChunk.BifrostError
+	}
+
+	// First chunk is valid data — wrap channel to re-inject it
+	done := make(chan struct{})
+	wrapped := make(chan *schemas.BifrostStreamChunk, max(cap(stream), 1))
+	wrapped <- firstChunk
+	go func() {
+		defer close(done)
+		defer close(wrapped)
+		for chunk := range stream {
+			select {
+			case wrapped <- chunk:
+			case <-ctx.Done():
+				// Consumer abandoned the wrapped channel. Drain the source so the
+				// provider's blocked send unblocks and its goroutine can exit.
+				for range stream {
+				}
+				return
+			}
+		}
+	}()
+	return wrapped, done, nil
+}