bugfix

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