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

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.
// convSvc must be the concrete *conversationService produced by NewConversationService;
// the type assertion lets callers in other packages pass the interface without needing
// to import the unexported concrete type. Pass nil to skip title generation.
func NewStreamerHub(
	repo Repository,
	registry clientGetter,
	dispatch notifyDispatcher,
	auditor audit.Recorder,
	convSvc ConversationService,
	log *slog.Logger,
	retention time.Duration,
	safetyPct float64,
) StreamerHub {
	var concrete *conversationService
	if convSvc != nil {
		if c, ok := convSvc.(*conversationService); ok {
			concrete = c
		}
	}
	return &streamerHub{
		byMsgID:   make(map[int64]*streamer),
		repo:      repo,
		registry:  registry,
		dispatch:  dispatch,
		auditor:   auditor,
		convSvc:   concrete,
		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
}