diff --git a/internal/chat/pricing.go b/internal/chat/pricing.go new file mode 100644 index 0000000..3923eb9 --- /dev/null +++ b/internal/chat/pricing.go @@ -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) } diff --git a/internal/chat/pricing_test.go b/internal/chat/pricing_test.go new file mode 100644 index 0000000..ed69a7f --- /dev/null +++ b/internal/chat/pricing_test.go @@ -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)) +} diff --git a/internal/chat/streamer.go b/internal/chat/streamer.go new file mode 100644 index 0000000..6d76f39 --- /dev/null +++ b/internal/chat/streamer.go @@ -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 +} diff --git a/internal/chat/streamer_test.go b/internal/chat/streamer_test.go new file mode 100644 index 0000000..1c78fe7 --- /dev/null +++ b/internal/chat/streamer_test.go @@ -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) diff --git a/internal/infra/llm/fake.go b/internal/infra/llm/fake.go index b793adc..697f685 100644 --- a/internal/infra/llm/fake.go +++ b/internal/infra/llm/fake.go @@ -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 }