feat(chat): StreamerHub with ring buffer + cond + Done/Error/Cancel flows

2026-05-04 17:27:50 +08:00
parent d1be9e6eb0
commit 617794c7a3
5 changed files with 1181 additions and 0 deletions
@@ -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
+}