bugfix

internal/chat/streamer.go

@@ -70,26 +70,50 @@ type subscription struct {
 // the stream closes, or ctx is cancelled.
 // It returns all events from nextEventID to the current tail.
 func (sub *subscription) Next(ctx context.Context, nextEventID int64) ([]SSEEvent, bool, error) {
-	sub.st.mu.Lock()
-	defer sub.st.mu.Unlock()
+	st := sub.st
 
-	for int64(len(sub.st.events)) <= nextEventID && !sub.st.closed {
-		done := make(chan struct{})
-		go func() {
-			select {
-			case <-ctx.Done():
-				sub.st.cond.Broadcast()
-			case <-done:
-			}
-		}()
-		sub.st.cond.Wait()
-		close(done)
-		if ctx.Err() != nil {
-			return nil, false, ctx.Err()
+	// Launch a single watcher that observes ctx cancellation and wakes any
+	// waiter. It broadcasts under the lock so the signal cannot be missed by a
+	// waiter that has not yet parked on cond.Wait(); the predicate loop below
+	// re-checks ctx.Err() after every wake, so a broadcast delivered before a
+	// waiter parks is still observed on the next predicate check.
+	//
+	// done is closed (under the lock) when Next returns, which makes the watcher
+	// exit even if ctx never cancels — so neither goroutine can leak.
+	done := make(chan struct{})
+	go func() {
+		select {
+		case <-ctx.Done():
+			st.mu.Lock()
+			st.cond.Broadcast()
+			st.mu.Unlock()
+		case <-done:
 		}
+	}()
+
+	st.mu.Lock()
+	defer func() {
+		// Signal the watcher to exit, then release the lock. Closing done under
+		// the lock pairs with the watcher's Broadcast-under-lock so there is no
+		// window in which the watcher broadcasts to an unlocked, already-returned
+		// Next.
+		close(done)
+		st.mu.Unlock()
+	}()
+
+	// Wait until there is data at/after nextEventID, the stream closes, or ctx
+	// is cancelled. Each predicate component is re-checked after every wake so a
+	// missed/early broadcast cannot wedge the wait.
+	for int64(len(st.events)) <= nextEventID && !st.closed && ctx.Err() == nil {
+		st.cond.Wait()
 	}
-	out := append([]SSEEvent(nil), sub.st.events[nextEventID:]...)
-	return out, sub.st.closed, nil
+
+	if int64(len(st.events)) <= nextEventID && !st.closed && ctx.Err() != nil {
+		return nil, false, ctx.Err()
+	}
+
+	out := append([]SSEEvent(nil), st.events[nextEventID:]...)
+	return out, st.closed, nil
 }
 
 // streamerHub manages all in-flight streamers.

internal/chat/streamer_test.go

@@ -176,7 +176,7 @@ func (panicRepo) SoftDeleteConversation(_ context.Context, _ uuid.UUID) error {
 func (panicRepo) InsertMessage(_ context.Context, _ uuid.UUID, _ MessageRole, _ string, _ MessageStatus) (*Message, error) {
 	panic("not implemented")
 }
-func (panicRepo) GetMessage(_ context.Context, _ int64) (*Message, error)  { panic("not implemented") }
+func (panicRepo) GetMessage(_ context.Context, _ int64) (*Message, error) { panic("not implemented") }
 func (panicRepo) ListMessagesByConversation(_ context.Context, _ uuid.UUID, _ *int64, _ int) ([]Message, error) {
 	panic("not implemented")
 }
@@ -667,6 +667,77 @@ func TestStreamer_TitleGenerationTriggered(t *testing.T) {
 	}, 1*time.Second, 10*time.Millisecond, "GenerateTitle should have been triggered")
 }
 
+// TestStreamer_Next_EarlyCancel_NoLeak exercises the early-cancel race in
+// subscription.Next: when ctx is already cancelled (or cancels in the window
+// before the waiter parks on cond.Wait) and there are no new events, Next must
+// return ctx.Err() promptly instead of blocking forever, and must not leak the
+// Next goroutine or its ctx watcher.
+func TestStreamer_Next_EarlyCancel_NoLeak(t *testing.T) {
+	t.Parallel()
+
+	st := newStreamer()
+	sub := &subscription{st: st}
+
+	// Case 1: ctx already cancelled before Next is called. With no events and an
+	// open stream, the predicate is unsatisfied except for ctx, so Next must
+	// observe cancellation immediately rather than parking forever.
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		evs, closed, err := sub.Next(ctx, 0)
+		require.ErrorIs(t, err, context.Canceled)
+		require.Nil(t, evs)
+		require.False(t, closed)
+	}()
+
+	select {
+	case <-done:
+	case <-time.After(2 * time.Second):
+		t.Fatal("Next did not return on an already-cancelled ctx; goroutine leaked")
+	}
+
+	// Case 2: cancel concurrently while a waiter is (about to be) parked. Repeat
+	// to widen the window between launching the watcher and parking on Wait.
+	for i := 0; i < 200; i++ {
+		ctx, cancel := context.WithCancel(context.Background())
+		ret := make(chan error, 1)
+		go func() {
+			_, _, err := sub.Next(ctx, 0)
+			ret <- err
+		}()
+		cancel()
+		select {
+		case err := <-ret:
+			require.ErrorIs(t, err, context.Canceled)
+		case <-time.After(2 * time.Second):
+			t.Fatalf("iteration %d: Next wedged on concurrent cancel", i)
+		}
+	}
+}
+
+// TestStreamer_Next_DataDeliveredDespiteCancel verifies that when events are
+// already buffered, Next returns them even if ctx is cancelled (data wins over
+// a late cancel, preserving the original behavior).
+func TestStreamer_Next_DataDeliveredDespiteCancel(t *testing.T) {
+	t.Parallel()
+
+	st := newStreamer()
+	sub := &subscription{st: st}
+	st.append(SSEEvent{Event: "text_delta", Data: map[string]string{"text": "hi"}})
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	evs, closed, err := sub.Next(ctx, 0)
+	require.NoError(t, err)
+	require.False(t, closed)
+	require.Len(t, evs, 1)
+	require.Equal(t, "text_delta", evs[0].Event)
+}
+
 // ===== helpers for title generation test =====
 
 // spyConvSvc is a minimal conversationService-shaped struct that calls a hook