You've already forked agentic-coding-workflow
726 lines
22 KiB
Go
726 lines
22 KiB
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) 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")
|
|
}
|
|
|
|
// ===== 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)
|