You've already forked agentic-coding-workflow
feat(chat): StreamerHub with ring buffer + cond + Done/Error/Cancel flows
This commit is contained in:
@@ -0,0 +1,15 @@
|
||||
package chat
|
||||
|
||||
import "strconv"
|
||||
|
||||
// computeCost returns the cost in USD for a single LLM call.
|
||||
// Prices on LLMModel are expressed in USD per million tokens.
|
||||
func computeCost(m LLMModel, ptok, ctok, ttok int) float64 {
|
||||
return (float64(ptok)*m.PromptPricePerMillionUSD +
|
||||
float64(ctok)*m.CompletionPricePerMillionUSD +
|
||||
float64(ttok)*m.ThinkingPricePerMillionUSD) / 1_000_000
|
||||
}
|
||||
|
||||
// int64ToString converts an int64 to its decimal string representation.
|
||||
// Uses strconv.FormatInt rather than fmt.Sprintf for efficiency.
|
||||
func int64ToString(v int64) string { return strconv.FormatInt(v, 10) }
|
||||
@@ -0,0 +1,47 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestComputeCost_BasicMath(t *testing.T) {
|
||||
m := LLMModel{
|
||||
PromptPricePerMillionUSD: 3.0,
|
||||
CompletionPricePerMillionUSD: 15.0,
|
||||
ThinkingPricePerMillionUSD: 3.0,
|
||||
}
|
||||
|
||||
// 1000 prompt tokens × $3/M = $0.003
|
||||
got := computeCost(m, 1000, 0, 0)
|
||||
require.InDelta(t, 0.003, got, 1e-9)
|
||||
|
||||
// 1000 completion tokens × $15/M = $0.015
|
||||
got = computeCost(m, 0, 1000, 0)
|
||||
require.InDelta(t, 0.015, got, 1e-9)
|
||||
|
||||
// 1000 thinking tokens × $3/M = $0.003
|
||||
got = computeCost(m, 0, 0, 1000)
|
||||
require.InDelta(t, 0.003, got, 1e-9)
|
||||
|
||||
// combined: 1000+1000+1000 = $0.021
|
||||
got = computeCost(m, 1000, 1000, 1000)
|
||||
require.InDelta(t, 0.021, got, 1e-9)
|
||||
}
|
||||
|
||||
func TestComputeCost_ZeroTokens(t *testing.T) {
|
||||
m := LLMModel{
|
||||
PromptPricePerMillionUSD: 3.0,
|
||||
CompletionPricePerMillionUSD: 15.0,
|
||||
ThinkingPricePerMillionUSD: 3.0,
|
||||
}
|
||||
require.Equal(t, 0.0, computeCost(m, 0, 0, 0))
|
||||
}
|
||||
|
||||
func TestInt64ToString(t *testing.T) {
|
||||
require.Equal(t, "0", int64ToString(0))
|
||||
require.Equal(t, "42", int64ToString(42))
|
||||
require.Equal(t, "-1", int64ToString(-1))
|
||||
require.Equal(t, "9223372036854775807", int64ToString(9223372036854775807))
|
||||
}
|
||||
@@ -0,0 +1,381 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"log/slog"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"github.com/yan1h/agent-coding-workflow/internal/audit"
|
||||
"github.com/yan1h/agent-coding-workflow/internal/infra/errs"
|
||||
"github.com/yan1h/agent-coding-workflow/internal/infra/llm"
|
||||
"github.com/yan1h/agent-coding-workflow/internal/infra/notify"
|
||||
)
|
||||
|
||||
// clientGetter is a testability seam around *llm.Registry.
|
||||
// *llm.Registry satisfies this interface implicitly.
|
||||
type clientGetter interface {
|
||||
GetClient(ctx context.Context, endpointID uuid.UUID) (llm.LLMClient, error)
|
||||
}
|
||||
|
||||
// notifyDispatcher is a testability seam around *notify.Dispatcher.
|
||||
// *notify.Dispatcher satisfies this interface implicitly.
|
||||
type notifyDispatcher interface {
|
||||
Dispatch(ctx context.Context, msg notify.Message) error
|
||||
}
|
||||
|
||||
// streamer holds the ring buffer and state for one in-flight assistant message.
|
||||
type streamer struct {
|
||||
mu sync.Mutex
|
||||
cond *sync.Cond
|
||||
events []SSEEvent
|
||||
closed bool
|
||||
cancel context.CancelFunc
|
||||
messageID int64
|
||||
}
|
||||
|
||||
func newStreamer() *streamer {
|
||||
s := &streamer{}
|
||||
s.cond = sync.NewCond(&s.mu)
|
||||
return s
|
||||
}
|
||||
|
||||
// append adds an event to the ring buffer and broadcasts to all waiters.
|
||||
// The event ID is assigned as the current length before append.
|
||||
func (s *streamer) append(ev SSEEvent) {
|
||||
s.mu.Lock()
|
||||
ev.ID = int64(len(s.events))
|
||||
s.events = append(s.events, ev)
|
||||
s.cond.Broadcast()
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
// close marks the streamer as done and wakes all waiters.
|
||||
func (s *streamer) close() {
|
||||
s.mu.Lock()
|
||||
s.closed = true
|
||||
s.cond.Broadcast()
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
// subscription is the handle returned to SSE HTTP handlers.
|
||||
type subscription struct {
|
||||
st *streamer
|
||||
}
|
||||
|
||||
// Next blocks until at least one event with ID >= lastEventID is available,
|
||||
// the stream closes, or ctx is cancelled.
|
||||
// It returns all events from lastEventID to the current tail.
|
||||
func (sub *subscription) Next(ctx context.Context, lastEventID int64) ([]SSEEvent, bool, error) {
|
||||
sub.st.mu.Lock()
|
||||
defer sub.st.mu.Unlock()
|
||||
|
||||
for int64(len(sub.st.events)) <= lastEventID && !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()
|
||||
}
|
||||
}
|
||||
out := append([]SSEEvent(nil), sub.st.events[lastEventID:]...)
|
||||
return out, sub.st.closed, nil
|
||||
}
|
||||
|
||||
// streamerHub manages all in-flight streamers.
|
||||
type streamerHub struct {
|
||||
mu sync.Mutex
|
||||
byMsgID map[int64]*streamer
|
||||
repo Repository
|
||||
registry clientGetter
|
||||
dispatch notifyDispatcher
|
||||
auditor audit.Recorder
|
||||
convSvc *conversationService
|
||||
log *slog.Logger
|
||||
retention time.Duration
|
||||
safety float64 // reserved for T15 token-truncation; unused in T14
|
||||
}
|
||||
|
||||
// NewStreamerHub constructs a StreamerHub.
|
||||
func NewStreamerHub(
|
||||
repo Repository,
|
||||
registry clientGetter,
|
||||
dispatch notifyDispatcher,
|
||||
auditor audit.Recorder,
|
||||
convSvc *conversationService,
|
||||
log *slog.Logger,
|
||||
retention time.Duration,
|
||||
safetyPct float64,
|
||||
) StreamerHub {
|
||||
return &streamerHub{
|
||||
byMsgID: make(map[int64]*streamer),
|
||||
repo: repo,
|
||||
registry: registry,
|
||||
dispatch: dispatch,
|
||||
auditor: auditor,
|
||||
convSvc: convSvc,
|
||||
log: log,
|
||||
retention: retention,
|
||||
safety: safetyPct,
|
||||
}
|
||||
}
|
||||
|
||||
// Start launches a goroutine to stream LLM output for the given params.
|
||||
// streamCtx is intentionally detached from the caller so the stream survives
|
||||
// after the initiating HTTP request completes.
|
||||
func (h *streamerHub) Start(ctx context.Context, p StreamParams) error {
|
||||
st := newStreamer()
|
||||
st.messageID = p.AssistantMessageID
|
||||
|
||||
streamCtx, cancel := context.WithCancel(context.Background())
|
||||
st.cancel = cancel
|
||||
|
||||
h.mu.Lock()
|
||||
h.byMsgID[p.AssistantMessageID] = st
|
||||
h.mu.Unlock()
|
||||
|
||||
go h.run(streamCtx, st, p)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Cancel signals the streaming goroutine to stop.
|
||||
// Returns true if the streamer was found, false otherwise.
|
||||
func (h *streamerHub) Cancel(messageID int64) bool {
|
||||
h.mu.Lock()
|
||||
st, ok := h.byMsgID[messageID]
|
||||
h.mu.Unlock()
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
st.cancel()
|
||||
return true
|
||||
}
|
||||
|
||||
// Subscribe returns a Subscription for the given messageID.
|
||||
// Returns false if no active or recently-completed streamer is found.
|
||||
func (h *streamerHub) Subscribe(messageID int64) (Subscription, bool) {
|
||||
h.mu.Lock()
|
||||
st, ok := h.byMsgID[messageID]
|
||||
h.mu.Unlock()
|
||||
if !ok {
|
||||
return nil, false
|
||||
}
|
||||
return &subscription{st: st}, true
|
||||
}
|
||||
|
||||
// run is the goroutine body for one stream. It calls the LLM, populates
|
||||
// the ring buffer, then persists the result in a transaction.
|
||||
func (h *streamerHub) run(ctx context.Context, st *streamer, p StreamParams) {
|
||||
defer h.scheduleRemoval(p.AssistantMessageID)
|
||||
defer st.close()
|
||||
|
||||
client, err := h.registry.GetClient(ctx, p.EndpointID)
|
||||
if err != nil {
|
||||
h.handleError(ctx, st, p, err)
|
||||
return
|
||||
}
|
||||
|
||||
st.append(SSEEvent{Event: "message_meta", Data: map[string]any{
|
||||
"id": p.AssistantMessageID,
|
||||
"role": "assistant",
|
||||
"status": "pending",
|
||||
}})
|
||||
|
||||
req := h.prepareRequest(ctx, p)
|
||||
|
||||
ch, err := client.Stream(ctx, p.ModelMeta.ModelID, req)
|
||||
if err != nil {
|
||||
h.handleError(ctx, st, p, err)
|
||||
return
|
||||
}
|
||||
|
||||
var content, thinking string
|
||||
var usage *llm.Usage
|
||||
var finishReason string
|
||||
|
||||
for ev := range ch {
|
||||
switch ev.Type {
|
||||
case llm.EventTextDelta:
|
||||
content += ev.Text
|
||||
st.append(SSEEvent{Event: "text_delta", Data: map[string]string{"text": ev.Text}})
|
||||
case llm.EventThinkingDelta:
|
||||
thinking += ev.Text
|
||||
st.append(SSEEvent{Event: "thinking_delta", Data: map[string]string{"text": ev.Text}})
|
||||
case llm.EventUsage:
|
||||
usage = ev.Usage
|
||||
st.append(SSEEvent{Event: "usage", Data: ev.Usage})
|
||||
case llm.EventDone:
|
||||
finishReason = ev.Reason
|
||||
case llm.EventError:
|
||||
h.handleError(ctx, st, p, ev.Err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// Cancel path: ctx.Err() != nil after channel drains.
|
||||
if errors.Is(ctx.Err(), context.Canceled) {
|
||||
_ = h.repo.UpdateMessageCancelled(context.Background(), p.AssistantMessageID)
|
||||
h.tryAudit(context.Background(), audit.Entry{
|
||||
UserID: &p.UserID,
|
||||
Action: "chat.message.cancelled",
|
||||
TargetType: "message",
|
||||
TargetID: int64ToString(p.AssistantMessageID),
|
||||
})
|
||||
st.append(SSEEvent{Event: "error", Data: map[string]string{"code": "cancelled"}})
|
||||
return
|
||||
}
|
||||
|
||||
// Done path: persist in a single transaction.
|
||||
bgCtx := context.Background()
|
||||
txRepo, tx, err := h.repo.WithTx(bgCtx)
|
||||
if err != nil {
|
||||
h.handleError(bgCtx, st, p, err)
|
||||
return
|
||||
}
|
||||
|
||||
committed := false
|
||||
defer func() {
|
||||
if !committed {
|
||||
_ = tx.Rollback(bgCtx)
|
||||
}
|
||||
}()
|
||||
|
||||
var thinkingPtr *string
|
||||
if thinking != "" {
|
||||
thinkingPtr = &thinking
|
||||
}
|
||||
|
||||
var ptok, ctok, ttok *int32
|
||||
if usage != nil {
|
||||
ptok = intToInt32Ptr(usage.PromptTokens)
|
||||
ctok = intToInt32Ptr(usage.CompletionTokens)
|
||||
ttok = intToInt32Ptr(usage.ThinkingTokens)
|
||||
}
|
||||
|
||||
if err := txRepo.UpdateMessageOK(bgCtx, UpdateMessageOKParams{
|
||||
ID: p.AssistantMessageID,
|
||||
Content: content,
|
||||
Thinking: thinkingPtr,
|
||||
PromptTokens: ptok,
|
||||
CompletionTokens: ctok,
|
||||
ThinkingTokens: ttok,
|
||||
}); err != nil {
|
||||
h.handleError(bgCtx, st, p, err)
|
||||
return
|
||||
}
|
||||
|
||||
if usage != nil {
|
||||
costUSD := computeCost(p.ModelMeta,
|
||||
usage.PromptTokens, usage.CompletionTokens, usage.ThinkingTokens)
|
||||
if err := txRepo.InsertUsage(bgCtx, UsageRecord{
|
||||
ConversationID: p.ConversationID,
|
||||
MessageID: p.AssistantMessageID,
|
||||
UserID: p.UserID,
|
||||
EndpointID: p.EndpointID,
|
||||
ModelID: p.ModelID,
|
||||
PromptTokens: usage.PromptTokens,
|
||||
CompletionTokens: usage.CompletionTokens,
|
||||
ThinkingTokens: usage.ThinkingTokens,
|
||||
CostUSD: costUSD,
|
||||
}); err != nil {
|
||||
h.handleError(bgCtx, st, p, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
if err := tx.Commit(bgCtx); err != nil {
|
||||
h.handleError(bgCtx, st, p, err)
|
||||
return
|
||||
}
|
||||
committed = true
|
||||
|
||||
h.tryAudit(bgCtx, audit.Entry{
|
||||
UserID: &p.UserID,
|
||||
Action: "chat.message.completed",
|
||||
TargetType: "message",
|
||||
TargetID: int64ToString(p.AssistantMessageID),
|
||||
Metadata: map[string]any{
|
||||
"finish_reason": finishReason,
|
||||
},
|
||||
})
|
||||
|
||||
st.append(SSEEvent{Event: "done", Data: map[string]string{"finish_reason": finishReason}})
|
||||
go h.maybeGenerateTitle(p.ConversationID)
|
||||
}
|
||||
|
||||
// handleError persists the error state, dispatches a user notification,
|
||||
// and appends an "error" SSE event to the stream.
|
||||
func (h *streamerHub) handleError(ctx context.Context, st *streamer, p StreamParams, err error) {
|
||||
bgCtx := context.Background()
|
||||
e := err.Error()
|
||||
_ = h.repo.UpdateMessageError(bgCtx, p.AssistantMessageID, &e)
|
||||
h.tryAudit(bgCtx, audit.Entry{
|
||||
UserID: &p.UserID,
|
||||
Action: "chat.message.error",
|
||||
TargetType: "message",
|
||||
TargetID: int64ToString(p.AssistantMessageID),
|
||||
Metadata: map[string]any{"error": e},
|
||||
})
|
||||
_ = h.dispatch.Dispatch(bgCtx, notify.Message{
|
||||
UserID: p.UserID,
|
||||
Topic: "chat.message_failed",
|
||||
Severity: notify.SeverityError,
|
||||
Title: "对话回复失败",
|
||||
Body: e,
|
||||
Link: "/chat/" + p.ConversationID.String(),
|
||||
})
|
||||
st.append(SSEEvent{Event: "error", Data: map[string]string{
|
||||
"code": string(errs.CodeLLMUpstream),
|
||||
"message": e,
|
||||
}})
|
||||
}
|
||||
|
||||
// scheduleRemoval removes the streamer from the hub after h.retention so
|
||||
// late-arriving subscribers can still catch the tail for a short window.
|
||||
func (h *streamerHub) scheduleRemoval(messageID int64) {
|
||||
time.AfterFunc(h.retention, func() {
|
||||
h.mu.Lock()
|
||||
delete(h.byMsgID, messageID)
|
||||
h.mu.Unlock()
|
||||
})
|
||||
}
|
||||
|
||||
// maybeGenerateTitle triggers title generation if the conversation has none yet.
|
||||
func (h *streamerHub) maybeGenerateTitle(conversationID uuid.UUID) {
|
||||
if h.convSvc == nil {
|
||||
return
|
||||
}
|
||||
c, err := h.repo.GetConversation(context.Background(), conversationID)
|
||||
if err != nil || c.Title != nil {
|
||||
return
|
||||
}
|
||||
h.convSvc.GenerateTitle(context.Background(), conversationID)
|
||||
}
|
||||
|
||||
// prepareRequest is a stub for T14; T15 will add context-window truncation.
|
||||
func (h *streamerHub) prepareRequest(_ context.Context, p StreamParams) llm.Request {
|
||||
return p.Request
|
||||
}
|
||||
|
||||
// tryAudit swallows audit errors, logging them at Warn level.
|
||||
func (h *streamerHub) tryAudit(ctx context.Context, e audit.Entry) {
|
||||
if err := h.auditor.Record(ctx, e); err != nil {
|
||||
h.log.Warn("audit record failed", "action", e.Action, "err", err)
|
||||
}
|
||||
}
|
||||
|
||||
// intToInt32Ptr converts an int to a heap-allocated int32 pointer.
|
||||
func intToInt32Ptr(v int) *int32 {
|
||||
x := int32(v)
|
||||
return &x
|
||||
}
|
||||
@@ -0,0 +1,725 @@
|
||||
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)
|
||||
@@ -15,6 +15,9 @@ type FakeClient struct {
|
||||
StreamError error // Stream() 立即返回的错误（如果非 nil）
|
||||
TokenCount int // CountTokens 返回值
|
||||
StreamCalls int
|
||||
// BlockCh, if non-nil, causes Stream goroutine to block after emitting all
|
||||
// Events until ctx is cancelled or BlockCh is closed. Useful for cancel tests.
|
||||
BlockCh chan struct{}
|
||||
}
|
||||
|
||||
// NewFake 构造一个 FakeClient，默认 provider="fake"。
|
||||
@@ -33,6 +36,8 @@ func (f *FakeClient) Stream(ctx context.Context, modelID string, req Request) (<
|
||||
events := append([]StreamEvent(nil), f.Events...)
|
||||
f.mu.Unlock()
|
||||
|
||||
blockCh := f.BlockCh
|
||||
|
||||
ch := make(chan StreamEvent, len(events)+1)
|
||||
go func() {
|
||||
defer close(ch)
|
||||
@@ -44,6 +49,14 @@ func (f *FakeClient) Stream(ctx context.Context, modelID string, req Request) (<
|
||||
case ch <- e:
|
||||
}
|
||||
}
|
||||
// If a block channel is set, pause here until ctx is cancelled or BlockCh
|
||||
// is closed. This lets cancel tests confirm mid-stream cancellation.
|
||||
if blockCh != nil {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
case <-blockCh:
|
||||
}
|
||||
}
|
||||
}()
|
||||
return ch, nil
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user