agentic-coding-workflow/internal/chat/streamer_test.go

package chat

import (
	"context"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/google/uuid"
	"github.com/stretchr/testify/require"

	"github.com/yan1h/agent-coding-workflow/internal/infra/llm"
	"github.com/yan1h/agent-coding-workflow/internal/infra/notify"
)

// ========== streamer-specific fakes ==========

// streamerFakeRepo is a test-only Repository whose relevant methods record calls
// instead of panicking. All unneeded methods are covered by the embedded
// panicRepo base.
type streamerFakeRepo struct {
	panicRepo // embed the panic-everything base defined below

	mu sync.Mutex

	conversations map[uuid.UUID]*Conversation
	messages      map[int64]*Message
	usageRecords  []UsageRecord

	updateOKCalled        bool
	updateOKParams        UpdateMessageOKParams
	updateErrorCalled     bool
	updateErrorID         int64
	updateErrorMsg        *string
	updateCancelledCalled bool
	updateCancelledID     int64
	insertUsageCalled     bool
	insertUsageRecord     UsageRecord
}

func newStreamerFakeRepo() *streamerFakeRepo {
	return &streamerFakeRepo{
		conversations: make(map[uuid.UUID]*Conversation),
		messages:      make(map[int64]*Message),
	}
}

func (r *streamerFakeRepo) GetConversation(_ context.Context, id uuid.UUID) (*Conversation, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	c, ok := r.conversations[id]
	if !ok {
		// Return a conversation with no title so title generation is triggered.
		return &Conversation{ID: id}, nil
	}
	return c, nil
}

func (r *streamerFakeRepo) GetTitleGeneratorModel(_ context.Context) (*LLMModel, error) {
	return nil, nil // no title model => GenerateTitle returns early
}

func (r *streamerFakeRepo) GetFirstUserMessage(_ context.Context, _ uuid.UUID) (*Message, error) {
	return nil, nil
}

func (r *streamerFakeRepo) UpdateMessageOK(_ context.Context, p UpdateMessageOKParams) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.updateOKCalled = true
	r.updateOKParams = p
	return nil
}

func (r *streamerFakeRepo) UpdateMessageError(_ context.Context, id int64, msg *string) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.updateErrorCalled = true
	r.updateErrorID = id
	r.updateErrorMsg = msg
	return nil
}

func (r *streamerFakeRepo) UpdateMessageCancelled(_ context.Context, id int64) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.updateCancelledCalled = true
	r.updateCancelledID = id
	return nil
}

func (r *streamerFakeRepo) InsertUsage(_ context.Context, rec UsageRecord) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.insertUsageCalled = true
	r.insertUsageRecord = rec
	r.usageRecords = append(r.usageRecords, rec)
	return nil
}

func (r *streamerFakeRepo) UpdateConversationTitle(_ context.Context, _ uuid.UUID, _ *string) error {
	return nil
}

// WithTx returns the same repo (mutations are visible) plus a no-op fakeTx.
func (r *streamerFakeRepo) WithTx(_ context.Context) (Repository, Tx, error) {
	return r, &fakeTx{}, nil
}

// ===== fakeTx =====

type fakeTx struct {
	committed  atomic.Bool
	rolledBack atomic.Bool
}

func (tx *fakeTx) Commit(_ context.Context) error {
	tx.committed.Store(true)
	return nil
}

func (tx *fakeTx) Rollback(_ context.Context) error {
	tx.rolledBack.Store(true)
	return nil
}

// ===== fakeClientGetter =====

type fakeClientGetter struct {
	client llm.LLMClient
}

func (g *fakeClientGetter) GetClient(_ context.Context, _ uuid.UUID) (llm.LLMClient, error) {
	return g.client, nil
}

// ===== fakeNotifyDispatcher =====

type fakeNotifyDispatcher struct {
	mu       sync.Mutex
	messages []notify.Message
}

func (d *fakeNotifyDispatcher) Dispatch(_ context.Context, msg notify.Message) error {
	d.mu.Lock()
	defer d.mu.Unlock()
	d.messages = append(d.messages, msg)
	return nil
}

func (d *fakeNotifyDispatcher) received() []notify.Message {
	d.mu.Lock()
	defer d.mu.Unlock()
	return append([]notify.Message(nil), d.messages...)
}

// ===== panicRepo =====

// panicRepo satisfies Repository and panics on every method.
// streamerFakeRepo embeds it and overrides only what it needs.
type panicRepo struct{}

func (panicRepo) InsertConversation(_ context.Context, _ InsertConversationParams) (*Conversation, error) {
	panic("not implemented")
}
func (panicRepo) ListConversationsByUser(_ context.Context, _ uuid.UUID, _ ConversationFilter) ([]Conversation, error) {
	panic("not implemented")
}
func (panicRepo) SetConversationTitleManual(_ context.Context, _ uuid.UUID, _ *string) error {
	panic("not implemented")
}
func (panicRepo) SoftDeleteConversation(_ context.Context, _ uuid.UUID) error {
	panic("not implemented")
}
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) ListMessagesByConversation(_ context.Context, _ uuid.UUID, _ *int64, _ int) ([]Message, error) {
	panic("not implemented")
}
func (panicRepo) ListMessagesAscending(_ context.Context, _ uuid.UUID) ([]Message, error) {
	panic("not implemented")
}
func (panicRepo) MarkPendingAsErrorOnStartup(_ context.Context) error { panic("not implemented") }
func (panicRepo) DeleteMessage(_ context.Context, _ int64) error      { panic("not implemented") }
func (panicRepo) InsertAttachment(_ context.Context, _ InsertAttachmentParams) (*Attachment, error) {
	panic("not implemented")
}
func (panicRepo) GetAttachment(_ context.Context, _ uuid.UUID) (*Attachment, error) {
	panic("not implemented")
}
func (panicRepo) AttachToMessage(_ context.Context, _ []uuid.UUID, _ int64) error {
	panic("not implemented")
}
func (panicRepo) ListAttachmentsForMessage(_ context.Context, _ int64) ([]Attachment, error) {
	panic("not implemented")
}
func (panicRepo) ListOrphanAttachmentsForUser(_ context.Context, _ uuid.UUID, _ []uuid.UUID) ([]Attachment, error) {
	panic("not implemented")
}
func (panicRepo) ListExpiredOrphans(_ context.Context, _ time.Duration) ([]AttachmentRef, error) {
	panic("not implemented")
}
func (panicRepo) ListAttachmentsForSoftDeletedConversations(_ context.Context, _ time.Duration) ([]AttachmentRef, error) {
	panic("not implemented")
}
func (panicRepo) DeleteAttachment(_ context.Context, _ uuid.UUID) error { panic("not implemented") }
func (panicRepo) DeleteAttachments(_ context.Context, _ []uuid.UUID) error {
	panic("not implemented")
}
func (panicRepo) InsertPromptTemplate(_ context.Context, _ InsertTemplateParams) (*PromptTemplate, error) {
	panic("not implemented")
}
func (panicRepo) GetPromptTemplate(_ context.Context, _ uuid.UUID) (*PromptTemplate, error) {
	panic("not implemented")
}
func (panicRepo) ListPromptTemplatesForUser(_ context.Context, _ uuid.UUID) ([]PromptTemplate, error) {
	panic("not implemented")
}
func (panicRepo) UpdatePromptTemplate(_ context.Context, _ uuid.UUID, _, _ string) error {
	panic("not implemented")
}
func (panicRepo) DeletePromptTemplate(_ context.Context, _ uuid.UUID) error {
	panic("not implemented")
}
func (panicRepo) InsertLLMEndpoint(_ context.Context, _ InsertEndpointParams) (*LLMEndpoint, error) {
	panic("not implemented")
}
func (panicRepo) GetLLMEndpoint(_ context.Context, _ uuid.UUID) (*LLMEndpoint, error) {
	panic("not implemented")
}
func (panicRepo) ListLLMEndpoints(_ context.Context) ([]LLMEndpoint, error) {
	panic("not implemented")
}
func (panicRepo) UpdateLLMEndpoint(_ context.Context, _ uuid.UUID, _ UpdateEndpointInput) error {
	panic("not implemented")
}
func (panicRepo) DeleteLLMEndpoint(_ context.Context, _ uuid.UUID) error { panic("not implemented") }
func (panicRepo) InsertLLMModel(_ context.Context, _ InsertModelParams) (*LLMModel, error) {
	panic("not implemented")
}
func (panicRepo) GetLLMModel(_ context.Context, _ uuid.UUID) (*LLMModel, error) {
	panic("not implemented")
}
func (panicRepo) ListLLMModels(_ context.Context, _ bool) ([]LLMModel, error) {
	panic("not implemented")
}
func (panicRepo) UpdateLLMModel(_ context.Context, _ uuid.UUID, _ UpdateModelInput) error {
	panic("not implemented")
}
func (panicRepo) DeleteLLMModel(_ context.Context, _ uuid.UUID) error { panic("not implemented") }
func (panicRepo) SummarizeUsageByUser(_ context.Context, _, _ time.Time) ([]UserUsageRow, error) {
	panic("not implemented")
}
func (panicRepo) SummarizeUsageByModel(_ context.Context, _, _ time.Time) ([]ModelUsageRow, error) {
	panic("not implemented")
}
func (panicRepo) SummarizeUsageByDay(_ context.Context, _, _ time.Time) ([]DayUsageRow, error) {
	panic("not implemented")
}
func (panicRepo) SummarizeUserDaily(_ context.Context, _ uuid.UUID, _, _ time.Time) ([]DayUsageRow, error) {
	panic("not implemented")
}
func (panicRepo) GetConversation(_ context.Context, _ uuid.UUID) (*Conversation, error) {
	panic("not implemented")
}
func (panicRepo) UpdateMessageOK(_ context.Context, _ UpdateMessageOKParams) error {
	panic("not implemented")
}
func (panicRepo) UpdateMessageError(_ context.Context, _ int64, _ *string) error {
	panic("not implemented")
}
func (panicRepo) UpdateMessageCancelled(_ context.Context, _ int64) error {
	panic("not implemented")
}
func (panicRepo) InsertUsage(_ context.Context, _ UsageRecord) error { panic("not implemented") }
func (panicRepo) WithTx(_ context.Context) (Repository, Tx, error)   { panic("not implemented") }
func (panicRepo) UpdateConversationTitle(_ context.Context, _ uuid.UUID, _ *string) error {
	panic("not implemented")
}
func (panicRepo) GetTitleGeneratorModel(_ context.Context) (*LLMModel, error) {
	panic("not implemented")
}
func (panicRepo) GetFirstUserMessage(_ context.Context, _ uuid.UUID) (*Message, error) {
	panic("not implemented")
}

// ===== test helpers =====

// buildStreamerHub builds a hub with controllable fakes.
func buildStreamerHub(
	t *testing.T,
	repo *streamerFakeRepo,
	fake *llm.FakeClient,
	dispatcher *fakeNotifyDispatcher,
) *streamerHub {
	t.Helper()
	getter := &fakeClientGetter{client: fake}
	hub := &streamerHub{
		byMsgID:   make(map[int64]*streamer),
		repo:      repo,
		registry:  getter,
		dispatch:  dispatcher,
		auditor:   noopAuditor{},
		convSvc:   nil, // title generation goes through maybeGenerateTitle -> early return
		log:       discardLogger(),
		retention: 50 * time.Millisecond,
		safety:    0.9,
	}
	return hub
}

// defaultParams returns a StreamParams with sensible defaults for testing.
func defaultParams(msgID int64) StreamParams {
	return StreamParams{
		UserID:             uuid.New(),
		ConversationID:     uuid.New(),
		EndpointID:         uuid.New(),
		ModelID:            uuid.New(),
		AssistantMessageID: msgID,
		ModelMeta: LLMModel{
			ModelID:                      "test-model",
			PromptPricePerMillionUSD:     3.0,
			CompletionPricePerMillionUSD: 15.0,
			ThinkingPricePerMillionUSD:   3.0,
		},
	}
}

// drainSub reads all events from a subscription until the stream closes or
// ctx expires. Returns all accumulated events.
func drainSub(ctx context.Context, sub Subscription) []SSEEvent {
	var all []SSEEvent
	var lastID int64
	for {
		evs, closed, err := sub.Next(ctx, lastID)
		if err != nil {
			return all
		}
		all = append(all, evs...)
		if len(evs) > 0 {
			lastID = evs[len(evs)-1].ID + 1
		}
		if closed {
			return all
		}
	}
}

// ===== Tests =====

func TestStreamer_DoneFlow_LandsContentAndUsage(t *testing.T) {
	t.Parallel()

	repo := newStreamerFakeRepo()
	fake := llm.NewFake()
	fake.Events = []llm.StreamEvent{
		{Type: llm.EventTextDelta, Text: "Hello"},
		{Type: llm.EventTextDelta, Text: " World"},
		{Type: llm.EventTextDelta, Text: "!"},
		{Type: llm.EventUsage, Usage: &llm.Usage{PromptTokens: 10, CompletionTokens: 20}},
		{Type: llm.EventDone, Reason: "end_turn"},
	}
	dispatcher := &fakeNotifyDispatcher{}
	hub := buildStreamerHub(t, repo, fake, dispatcher)

	msgID := int64(42)
	p := defaultParams(msgID)
	require.NoError(t, hub.Start(context.Background(), p))

	sub, ok := hub.Subscribe(msgID)
	require.True(t, ok)

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	evs := drainSub(ctx, sub)

	require.Eventually(t, func() bool {
		repo.mu.Lock()
		defer repo.mu.Unlock()
		return repo.updateOKCalled
	}, 1*time.Second, 10*time.Millisecond, "UpdateMessageOK should have been called")

	repo.mu.Lock()
	params := repo.updateOKParams
	repo.mu.Unlock()

	require.Equal(t, msgID, params.ID)
	require.Equal(t, "Hello World!", params.Content)

	require.Eventually(t, func() bool {
		repo.mu.Lock()
		defer repo.mu.Unlock()
		return repo.insertUsageCalled
	}, 1*time.Second, 10*time.Millisecond, "InsertUsage should have been called")

	repo.mu.Lock()
	rec := repo.insertUsageRecord
	repo.mu.Unlock()

	require.Equal(t, 10, rec.PromptTokens)
	require.Equal(t, 20, rec.CompletionTokens)

	// Verify SSE events contain text_delta and done.
	var eventTypes []string
	for _, e := range evs {
		eventTypes = append(eventTypes, e.Event)
	}
	require.Contains(t, eventTypes, "text_delta")
	require.Contains(t, eventTypes, "done")
}

func TestStreamer_ErrorFlow_DispatchesNotify(t *testing.T) {
	t.Parallel()

	repo := newStreamerFakeRepo()
	fake := llm.NewFake()
	fake.Events = []llm.StreamEvent{
		{Type: llm.EventError, Err: &testError{"upstream exploded"}},
	}
	dispatcher := &fakeNotifyDispatcher{}
	hub := buildStreamerHub(t, repo, fake, dispatcher)

	msgID := int64(99)
	p := defaultParams(msgID)
	require.NoError(t, hub.Start(context.Background(), p))

	sub, ok := hub.Subscribe(msgID)
	require.True(t, ok)

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	evs := drainSub(ctx, sub)

	// UpdateMessageError must have been called.
	require.Eventually(t, func() bool {
		repo.mu.Lock()
		defer repo.mu.Unlock()
		return repo.updateErrorCalled
	}, 1*time.Second, 10*time.Millisecond)

	repo.mu.Lock()
	errID := repo.updateErrorID
	errMsg := repo.updateErrorMsg
	repo.mu.Unlock()

	require.Equal(t, msgID, errID)
	require.NotNil(t, errMsg)
	require.Contains(t, *errMsg, "upstream exploded")

	// Dispatcher must have received a notification.
	require.Eventually(t, func() bool {
		return len(dispatcher.received()) > 0
	}, 1*time.Second, 10*time.Millisecond)

	msgs := dispatcher.received()
	require.Len(t, msgs, 1)
	require.Equal(t, "chat.message_failed", msgs[0].Topic)

	// Last SSE event must be "error" with code "llm.upstream_error".
	var lastEvent SSEEvent
	for _, e := range evs {
		lastEvent = e
	}
	require.Equal(t, "error", lastEvent.Event)
	dataMap, ok := lastEvent.Data.(map[string]string)
	require.True(t, ok)
	require.Equal(t, "llm.upstream_error", dataMap["code"])
}

func TestStreamer_CancelMidStream_StatusCancelled(t *testing.T) {
	t.Parallel()

	repo := newStreamerFakeRepo()
	fake := llm.NewFake()
	blockCh := make(chan struct{})
	fake.BlockCh = blockCh
	fake.Events = []llm.StreamEvent{
		{Type: llm.EventTextDelta, Text: "partial"},
	}
	dispatcher := &fakeNotifyDispatcher{}
	hub := buildStreamerHub(t, repo, fake, dispatcher)

	msgID := int64(7)
	p := defaultParams(msgID)
	require.NoError(t, hub.Start(context.Background(), p))

	// Subscribe before cancelling so we can read the final error event.
	sub, ok := hub.Subscribe(msgID)
	require.True(t, ok)

	// Give the goroutine a moment to emit the partial event and reach the block.
	require.Eventually(t, func() bool {
		// peek: has at least the text_delta (ID=1, after message_meta ID=0)
		evs, _, _ := sub.Next(context.Background(), 0)
		for _, e := range evs {
			if e.Event == "text_delta" {
				return true
			}
		}
		return false
	}, 1*time.Second, 5*time.Millisecond)

	hub.Cancel(msgID)

	require.Eventually(t, func() bool {
		repo.mu.Lock()
		defer repo.mu.Unlock()
		return repo.updateCancelledCalled
	}, 1*time.Second, 10*time.Millisecond, "UpdateMessageCancelled should have been called")

	require.Equal(t, msgID, repo.updateCancelledID)

	// Drain remaining events and find the "error" cancelled event.
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	evs := drainSub(ctx, sub)

	var hasCancel bool
	for _, e := range evs {
		if e.Event == "error" {
			dataMap, ok := e.Data.(map[string]string)
			if ok && dataMap["code"] == "cancelled" {
				hasCancel = true
			}
		}
	}
	require.True(t, hasCancel, "expected SSE error event with code=cancelled")
}

func TestStreamer_Subscribe_ReceivesEventsInOrder(t *testing.T) {
	t.Parallel()

	repo := newStreamerFakeRepo()
	fake := llm.NewFake()
	fake.Events = []llm.StreamEvent{
		{Type: llm.EventTextDelta, Text: "a"},
		{Type: llm.EventTextDelta, Text: "b"},
		{Type: llm.EventTextDelta, Text: "c"},
		{Type: llm.EventDone, Reason: "end_turn"},
	}
	dispatcher := &fakeNotifyDispatcher{}
	hub := buildStreamerHub(t, repo, fake, dispatcher)

	msgID := int64(5)
	p := defaultParams(msgID)
	require.NoError(t, hub.Start(context.Background(), p))

	sub, ok := hub.Subscribe(msgID)
	require.True(t, ok)

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	evs := drainSub(ctx, sub)

	// IDs must be monotonically increasing.
	for i := 1; i < len(evs); i++ {
		require.Greater(t, evs[i].ID, evs[i-1].ID,
			"event IDs must be strictly increasing; got %v then %v", evs[i-1].ID, evs[i].ID)
	}

	// Events must contain the three text_delta values in sequence.
	var deltas []string
	for _, e := range evs {
		if e.Event == "text_delta" {
			d, ok := e.Data.(map[string]string)
			require.True(t, ok)
			deltas = append(deltas, d["text"])
		}
	}
	require.Equal(t, []string{"a", "b", "c"}, deltas)
}

func TestStreamer_RetentionRemovesAfterTimer(t *testing.T) {
	t.Parallel()

	repo := newStreamerFakeRepo()
	fake := llm.NewFake()
	fake.Events = []llm.StreamEvent{
		{Type: llm.EventDone, Reason: "end_turn"},
	}
	dispatcher := &fakeNotifyDispatcher{}

	hub := &streamerHub{
		byMsgID:   make(map[int64]*streamer),
		repo:      repo,
		registry:  &fakeClientGetter{client: fake},
		dispatch:  dispatcher,
		auditor:   noopAuditor{},
		convSvc:   nil,
		log:       discardLogger(),
		retention: 20 * time.Millisecond, // very short for test speed
		safety:    0.9,
	}

	msgID := int64(3)
	p := defaultParams(msgID)
	require.NoError(t, hub.Start(context.Background(), p))

	// Stream should complete; wait for it.
	sub, ok := hub.Subscribe(msgID)
	require.True(t, ok)
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	drainSub(ctx, sub)

	// After retention elapses, the entry should be gone.
	require.Eventually(t, func() bool {
		hub.mu.Lock()
		defer hub.mu.Unlock()
		_, exists := hub.byMsgID[msgID]
		return !exists
	}, 500*time.Millisecond, 5*time.Millisecond, "byMsgID entry should be removed after retention")
}

func TestStreamer_TitleGenerationTriggered(t *testing.T) {
	t.Parallel()

	repo := newStreamerFakeRepo()
	convoID := uuid.New()

	// Set up a fake conversation with no title so GenerateTitle path runs.
	repo.conversations[convoID] = &Conversation{ID: convoID, Title: nil}

	fake := llm.NewFake()
	fake.Events = []llm.StreamEvent{
		{Type: llm.EventTextDelta, Text: "hello"},
		{Type: llm.EventDone, Reason: "end_turn"},
	}
	dispatcher := &fakeNotifyDispatcher{}

	// We need a real conversationService-like thing to spy on.
	// Use a titleTriggeredFlag that we can set via a custom convSvc wrapper.
	var titleTriggered atomic.Bool
	convSvc := &spyConvSvc{onGenerateTitle: func() { titleTriggered.Store(true) }}

	hub := &streamerHub{
		byMsgID:   make(map[int64]*streamer),
		repo:      repo,
		registry:  &fakeClientGetter{client: fake},
		dispatch:  dispatcher,
		auditor:   noopAuditor{},
		convSvc:   convSvc.asConversationService(),
		log:       discardLogger(),
		retention: 50 * time.Millisecond,
		safety:    0.9,
	}

	msgID := int64(11)
	p := defaultParams(msgID)
	p.ConversationID = convoID
	require.NoError(t, hub.Start(context.Background(), p))

	sub, ok := hub.Subscribe(msgID)
	require.True(t, ok)
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	drainSub(ctx, sub)

	require.Eventually(t, func() bool {
		return titleTriggered.Load()
	}, 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
// when GenerateTitle is invoked. It wraps the concrete type via embedding.
type spyConvSvc struct {
	onGenerateTitle func()
}

// asConversationService returns a *conversationService whose GenerateTitle
// method we intercept by using a specially crafted repo.
func (s *spyConvSvc) asConversationService() *conversationService {
	// We can't monkey-patch the method, so we use a dedicated spy repo instead.
	// See spyConvSvcRepo below.
	return &conversationService{
		repo:    &spyConvSvcRepo{hook: s.onGenerateTitle},
		log:     discardLogger(),
		auditor: noopAuditor{},
	}
}

// spyConvSvcRepo is a repository whose GetTitleGeneratorModel returns a fake
// model, causing GenerateTitle to attempt a stream — but we override
// GetFirstUserMessage to return an error so it bails early, and we call the
// hook right before the bail.
type spyConvSvcRepo struct {
	panicRepo
	hook func()
}

func (r *spyConvSvcRepo) GetConversation(_ context.Context, id uuid.UUID) (*Conversation, error) {
	return &Conversation{ID: id, Title: nil}, nil
}

func (r *spyConvSvcRepo) GetTitleGeneratorModel(_ context.Context) (*LLMModel, error) {
	r.hook() // signal that title generation was triggered
	return nil, &testError{"no model"}
}

func (r *spyConvSvcRepo) GetFirstUserMessage(_ context.Context, _ uuid.UUID) (*Message, error) {
	panic("should not reach here")
}

func (r *spyConvSvcRepo) UpdateConversationTitle(_ context.Context, _ uuid.UUID, _ *string) error {
	return nil
}

// testError is a simple error type for tests.
type testError struct{ msg string }

func (e *testError) Error() string { return e.msg }

// Compile-time assertion: streamerFakeRepo satisfies Repository.
var _ Repository = (*streamerFakeRepo)(nil)

// Compile-time assertion: panicRepo satisfies Repository.
var _ Repository = (*panicRepo)(nil)

// Compile-time assertion: spyConvSvcRepo satisfies Repository.
var _ Repository = (*spyConvSvcRepo)(nil)