Fix Openstream message type

Author: hkotian

internal/dev_server/api/events/events_stream.go

@@ -16,10 +16,10 @@ import (
 type sdkEventObserver struct {
 	ctx             context.Context
 	debugSessionKey string
-	updateChan      chan<- sdk.Message
+	updateChan      chan<- []byte
 }
 
-func newSdkEventObserver(updateChan chan<- sdk.Message, ctx context.Context) sdkEventObserver {
+func newSdkEventObserver(updateChan chan<- []byte, ctx context.Context) sdkEventObserver {
 	debugSessionKey := uuid.New().String()
 	db := model.EventStoreFromContext(ctx)
 	err := db.CreateDebugSession(ctx, debugSessionKey)
@@ -54,14 +54,14 @@ func (o sdkEventObserver) Handle(message interface{}) {
 		return
 	}
 
-	o.updateChan <- sdk.Message{Event: sdk.TYPE_PUT, Data: str}
+	o.updateChan <- sdk.Message{Event: sdk.TYPE_PUT, Data: str}.ToPayload()
 }
 
 func SdkEventsTeeHandler(writer http.ResponseWriter, request *http.Request) {
 	updateChan, errChan := sdk.OpenStream(
 		writer,
 		request.Context().Done(),
-		sdk.Message{Event: sdk.TYPE_PUT, Data: []byte{}},
+		sdk.Message{Event: sdk.TYPE_PUT, Data: []byte{}}.ToPayload(),
 	)
 	defer close(updateChan)
 	observers := model.GetObserversFromContext(request.Context())

internal/dev_server/sdk/fdv2.go

@@ -123,7 +123,7 @@ func makePutObjectEvent(version int, key string, flagState model.FlagState) (sub
 	return subsystems.RawEvent{Name: subsystems.EventPutObject, Data: data}, nil
 }
 
-// buildFlagChangeEvents builds the three-event sequence for a single flag update pushed over a stream:
+// buildFlagChangeEvents builds the events sequence for a single flag update pushed over a stream:
 // server-intent(xfer-changes) + put-object(changed flag) + payload-transferred.
 func buildFlagChangeEvents(payloadID string, version int, flagKey string, flagState model.FlagState) ([]subsystems.RawEvent, error) {
 	intentEvent, err := makeServerIntentEvent(payloadID, version, subsystems.IntentTransferChanges, fdv2ReasonUpdate)

internal/dev_server/sdk/stream_client_flags.go

@@ -25,7 +25,7 @@ func StreamClientFlags(w http.ResponseWriter, r *http.Request) {
 	updateChan, doneChan := OpenStream(
 		w,
 		r.Context().Done(),
-		Message{Event: TYPE_PUT, Data: jsonBody},
+		Message{Event: TYPE_PUT, Data: jsonBody}.ToPayload(),
 	)
 	defer close(updateChan)
 	projectKey := GetProjectKeyFromContext(ctx)
@@ -46,7 +46,7 @@ func StreamClientFlags(w http.ResponseWriter, r *http.Request) {
 }
 
 type clientFlagsObserver struct {
-	updateChan chan<- Message
+	updateChan chan<- []byte
 	projectKey string
 }
 

internal/dev_server/sdk/stream_server_fdv2.go

@@ -4,7 +4,6 @@ import (
 	"fmt"
 	"log"
 	"net/http"
-	"time"
 
 	"github.com/launchdarkly/go-server-sdk/v7/subsystems"
 	"github.com/launchdarkly/ldcli/internal/dev_server/model"
@@ -28,72 +27,41 @@ func StreamV2(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	flusher, ok := w.(http.Flusher)
-	if !ok {
-		WriteError(ctx, w, errors.New("streaming not supported"))
-		return
-	}
-
 	initialPayload, err := buildFullTransferResponse(projectKey, project.PayloadVersion, allFlags, fdv2ReasonPayloadMissing)
 	if err != nil {
 		WriteError(ctx, w, errors.Wrap(err, "failed to build initial payload"))
 		return
 	}
 
-	// Register observer before writing to the client so that any changes arriving
-	// during the initial write are queued and delivered immediately after.
-	updateChan := make(chan []subsystems.RawEvent, 10)
-	observerID := model.GetObserversFromContext(ctx).RegisterObserver(fdv2StreamObserver{
-		updateChan: updateChan,
-		projectKey: projectKey,
-	})
+	updateChan, doneChan := OpenStream(w, r.Context().Done(), fdv2SSEPayload(initialPayload.Events))
+	defer close(updateChan)
+
+	observer := fdv2StreamObserver{updateChan: updateChan, projectKey: projectKey}
+	observerID := model.GetObserversFromContext(ctx).RegisterObserver(observer)
 	defer func() {
 		if ok := model.GetObserversFromContext(ctx).DeregisterObserver(observerID); !ok {
 			log.Printf("unable to deregister fdv2 stream observer")
 		}
 	}()
 
-	w.Header().Set("Content-Type", "text/event-stream")
-	w.Header().Set("Cache-Control", "no-cache")
-
-	if err := writeFDv2SSEEvents(w, flusher, initialPayload.Events); err != nil {
-		return
-	}
-
-	ticker := time.NewTicker(time.Minute)
-	defer ticker.Stop()
-
-	for {
-		select {
-		case events := <-updateChan:
-			if err := writeFDv2SSEEvents(w, flusher, events); err != nil {
-				return
-			}
-		case <-ticker.C:
-			// SSE comment line as a keepalive.
-			if _, err := w.Write([]byte(":\n\n")); err != nil {
-				return
-			}
-			flusher.Flush()
-		case <-ctx.Done():
-			return
-		}
+	err = <-doneChan
+	if err != nil {
+		WriteError(ctx, w, errors.Wrap(err, "stream failure"))
 	}
 }
 
-// writeFDv2SSEEvents writes a batch of FDv2 events to the response as individual SSE events.
-func writeFDv2SSEEvents(w http.ResponseWriter, flusher http.Flusher, events []subsystems.RawEvent) error {
-	for _, event := range events {
-		if _, err := fmt.Fprintf(w, "event:%s\ndata:%s\n\n", event.Name, event.Data); err != nil {
-			return err
-		}
+// fdv2SSEPayload formats a slice of FDv2 events as raw SSE bytes.
+// Each event becomes an individual SSE event in the output.
+func fdv2SSEPayload(events []subsystems.RawEvent) []byte {
+	var buf []byte
+	for _, e := range events {
+		buf = append(buf, fmt.Sprintf("event:%s\ndata:%s\n\n", e.Name, e.Data)...)
 	}
-	flusher.Flush()
-	return nil
+	return buf
 }
 
 type fdv2StreamObserver struct {
-	updateChan chan<- []subsystems.RawEvent
+	updateChan chan<- []byte
 	projectKey string
 }
 
@@ -107,7 +75,7 @@ func (o fdv2StreamObserver) Handle(event interface{}) {
 		if err != nil {
 			panic(errors.Wrap(err, "failed to build flag change events in fdv2 stream observer"))
 		}
-		o.updateChan <- events
+		o.updateChan <- fdv2SSEPayload(events)
 	case model.SyncEvent:
 		if event.ProjectKey != o.projectKey {
 			return
@@ -116,6 +84,6 @@ func (o fdv2StreamObserver) Handle(event interface{}) {
 		if err != nil {
 			panic(errors.Wrap(err, "failed to build full transfer in fdv2 stream observer"))
 		}
-		o.updateChan <- payload.Events
+		o.updateChan <- fdv2SSEPayload(payload.Events)
 	}
 }

internal/dev_server/sdk/stream_server_flags.go

@@ -27,7 +27,7 @@ func StreamServerAllPayload(w http.ResponseWriter, r *http.Request) {
 	updateChan, doneChan := OpenStream(
 		w,
 		r.Context().Done(),
-		Message{Event: TYPE_PUT, Data: jsonBody},
+		Message{Event: TYPE_PUT, Data: jsonBody}.ToPayload(),
 	)
 	defer close(updateChan)
 	observer := serverFlagsObserver{updateChan, projectKey}
@@ -47,7 +47,7 @@ func StreamServerAllPayload(w http.ResponseWriter, r *http.Request) {
 }
 
 type serverFlagsObserver struct {
-	updateChan chan<- Message
+	updateChan chan<- []byte
 	projectKey string
 }
 

internal/dev_server/sdk/streaming.go

@@ -28,10 +28,11 @@ func (m Message) ToPayload() []byte {
 	return payload
 }
 
-// OpenStream sends data to a response using the initial payload and subsequently via the returned write only channel
-func OpenStream(w http.ResponseWriter, done <-chan struct{}, initialMessage Message) (chan<- Message, <-chan error) {
+// OpenStream sets SSE headers, writes initialPayload, and starts the SSE loop.
+// Each []byte sent to the returned channel is written verbatim to the response.
+func OpenStream(w http.ResponseWriter, done <-chan struct{}, initialPayload []byte) (chan<- []byte, <-chan error) {
 	errChan := make(chan error)
-	updateChan := make(chan Message, 10)
+	updateChan := make(chan []byte, 10)
 	go func() {
 		var err error
 		defer func() {
@@ -45,7 +46,7 @@ func OpenStream(w http.ResponseWriter, done <-chan struct{}, initialMessage Mess
 			}
 
 			w.Header().Set("Content-Type", "text/event-stream")
-			_, err = w.Write(initialMessage.ToPayload())
+			_, err = w.Write(initialPayload)
 			if err != nil {
 				return errors.Wrap(err, "unable to write response")
 			}
@@ -60,8 +61,8 @@ func OpenStream(w http.ResponseWriter, done <-chan struct{}, initialMessage Mess
 						return errors.Wrap(err, "unable to write response")
 					}
 					flusher.Flush()
-				case msg := <-updateChan:
-					_, err = w.Write(msg.ToPayload())
+				case payload := <-updateChan:
+					_, err = w.Write(payload)
 					if err != nil {
 						return errors.Wrap(err, "unable to write response")
 					}
@@ -77,7 +78,7 @@ func OpenStream(w http.ResponseWriter, done <-chan struct{}, initialMessage Mess
 }
 
 func SendMessage(
-	updateChan chan<- Message,
+	updateChan chan<- []byte,
 	msgType MessageType,
 	data interface{},
 ) error {
@@ -89,7 +90,7 @@ func SendMessage(
 	updateChan <- Message{
 		Event: msgType,
 		Data:  payload,
-	}
+	}.ToPayload()
 
 	return nil
 }