[Backport 5.3] Fix panics in streaming completions handler (#60476)

Fix panics in streaming completions handler (#60474)

(cherry picked from commit 52dffbb)

Co-authored-by: Cezary Bartoszuk <[email protected]>

Author: sourcegraph-release-bot

internal/completions/httpapi/handler.go

@@ -226,31 +226,34 @@ func newStreamingResponseHandler(logger log.Logger, db database.DB, feature type
 			return eventWriter
 		})
 
+		// Isolate writing events.
+		var mu sync.Mutex
+		writeEvent := func(name string, data any) error {
+			mu.Lock()
+			defer mu.Unlock()
+			if ev := eventWriter(); ev != nil {
+				return ev.Event(name, data)
+			}
+			return nil
+		}
+
 		// Always send a final done event so clients know the stream is shutting down.
 		firstEventObserved := false
 		defer func() {
 			if firstEventObserved {
-				if ev := eventWriter(); ev != nil {
-					_ = ev.Event("done", map[string]any{})
-				}
+				_ = writeEvent("done", map[string]any{})
 			}
 		}()
 		start := time.Now()
 		eventSink := func(e types.CompletionResponse) error {
-			if w := eventWriter(); w != nil {
-				return w.Event("completion", e)
-			}
-			return nil
+			return writeEvent("completion", e)
 		}
 		attributionErrorLog := func(err error) {
 			l := trace.Logger(ctx, logger)
-			ev := eventWriter()
-			if ev != nil {
-				if err := ev.Event("attribution-error", map[string]string{"error": err.Error()}); err != nil {
-					l.Error("error reporting attribution error", log.Error(err))
-				} else {
-					return
-				}
+			if err := writeEvent("attribution-error", map[string]string{"error": err.Error()}); err != nil {
+				l.Error("error reporting attribution error", log.Error(err))
+			} else {
+				return
 			}
 			l.Error("attribution error", log.Error(err))
 		}
@@ -325,20 +328,16 @@ func newStreamingResponseHandler(logger log.Logger, db database.DB, feature type
 				firstEventObserved = true
 				timeToFirstEventMetrics.Observe(time.Since(start).Seconds(), 1, nil, requestParams.Model)
 			}
-			if ev := eventWriter(); ev != nil {
-				if err := ev.Event("error", map[string]string{"error": err.Error()}); err != nil {
-					l.Error("error reporting streaming completion error", log.Error(err))
-				}
+			if err := writeEvent("error", map[string]string{"error": err.Error()}); err != nil {
+				l.Error("error reporting streaming completion error", log.Error(err))
 			}
 			return
 		}
 		if f != nil { // if autocomplete-attribution enabled
 			if err := f.WaitDone(ctx); err != nil {
 				l := trace.Logger(ctx, logger)
-				if ev := eventWriter(); ev != nil {
-					if err := ev.Event("error", map[string]string{"error": err.Error()}); err != nil {
-						l.Error("error reporting streaming completion error", log.Error(err))
-					}
+				if err := writeEvent("error", map[string]string{"error": err.Error()}); err != nil {
+					l.Error("error reporting streaming completion error", log.Error(err))
 				}
 			}
 		}

internal/guardrails/attribution_filter.go

@@ -115,6 +115,10 @@ func (a *completionsFilter) Send(ctx context.Context, e types.CompletionResponse
 // is called and it will wait for remaining attribution search if context
 // time limits permit.
 func (a *completionsFilter) WaitDone(ctx context.Context) error {
+	// If attribution never run, we're done.
+	a.attributionRun.Do(func() {
+		close(a.attributionFinished)
+	})
 	select {
 	case <-ctx.Done():
 		a.blockSending()