feat: stream persistent

2026-04-28 22:41:03 -04:00
parent fad8ed865a
commit eb66801f58
7 changed files with 485 additions and 133 deletions
--- a/backend/internal/api/generation.go
+++ b/backend/internal/api/generation.go
@@ -0,0 +1,194 @@
+package api
+
+import (
+	"errors"
+	"sync"
+
+	"github.com/google/uuid"
+	"reichard.io/aethera/internal/store"
+	"reichard.io/aethera/internal/types"
+)
+
+var errGenerationActive = errors.New("generation already active")
+
+type generationManager struct {
+	mu          sync.RWMutex
+	generations map[uuid.UUID]*generation
+}
+
+type generation struct {
+	mu          sync.RWMutex
+	subscribers map[chan *MessageChunk]struct{}
+	done        chan struct{}
+	closed      bool
+}
+
+func newGenerationManager() *generationManager {
+	return &generationManager{generations: make(map[uuid.UUID]*generation)}
+}
+
+func (m *generationManager) start(chatID uuid.UUID, prepare func(*generation) error, run func(*generation)) error {
+	m.mu.Lock()
+	if _, found := m.generations[chatID]; found {
+		m.mu.Unlock()
+		return errGenerationActive
+	}
+
+	// Reserve Generation
+	gen := &generation{
+		subscribers: make(map[chan *MessageChunk]struct{}),
+		done:        make(chan struct{}),
+	}
+	m.generations[chatID] = gen
+	m.mu.Unlock()
+
+	// Prepare Generation - This runs while the generation is reserved so a
+	// concurrent request cannot persist duplicate user/assistant messages.
+	if err := prepare(gen); err != nil {
+		gen.close()
+		m.mu.Lock()
+		delete(m.generations, chatID)
+		m.mu.Unlock()
+		return err
+	}
+
+	// Run Generation
+	go func() {
+		defer func() {
+			gen.close()
+			m.mu.Lock()
+			delete(m.generations, chatID)
+			m.mu.Unlock()
+		}()
+		run(gen)
+	}()
+
+	return nil
+}
+
+func (m *generationManager) subscribe(chatID uuid.UUID) (<-chan *MessageChunk, func(), bool) {
+	m.mu.RLock()
+	gen, found := m.generations[chatID]
+	m.mu.RUnlock()
+	if !found {
+		return nil, func() {}, false
+	}
+
+	ch := gen.subscribe()
+	return ch, func() { gen.unsubscribe(ch) }, true
+}
+
+func (g *generation) subscribe() chan *MessageChunk {
+	ch := make(chan *MessageChunk, 64)
+
+	// Add Subscriber
+	g.mu.Lock()
+	if g.closed {
+		close(ch)
+	} else {
+		g.subscribers[ch] = struct{}{}
+	}
+	g.mu.Unlock()
+
+	return ch
+}
+
+func (g *generation) unsubscribe(ch chan *MessageChunk) {
+	// Remove Subscriber
+	g.mu.Lock()
+	if _, found := g.subscribers[ch]; found {
+		delete(g.subscribers, ch)
+		close(ch)
+	}
+	g.mu.Unlock()
+}
+
+func (g *generation) broadcast(chunk *MessageChunk) {
+	g.mu.RLock()
+	defer g.mu.RUnlock()
+
+	// Broadcast Chunk
+	for subscriber := range g.subscribers {
+		select {
+		case subscriber <- cloneMessageChunk(chunk):
+		default:
+		}
+	}
+}
+
+func (g *generation) close() {
+	g.mu.Lock()
+	defer g.mu.Unlock()
+
+	// Close Subscribers
+	if g.closed {
+		return
+	}
+	g.closed = true
+	close(g.done)
+	for subscriber := range g.subscribers {
+		close(subscriber)
+		delete(g.subscribers, subscriber)
+	}
+}
+
+func cloneMessageChunk(chunk *MessageChunk) *MessageChunk {
+	if chunk == nil {
+		return nil
+	}
+
+	// Clone Chunk
+	cloned := &MessageChunk{
+		Chat:             chunk.Chat,
+		UserMessage:      cloneStoreMessage(chunk.UserMessage),
+		AssistantMessage: cloneStoreMessage(chunk.AssistantMessage),
+	}
+	return cloned
+}
+
+func cloneStoreMessage(msg *store.Message) *store.Message {
+	if msg == nil {
+		return nil
+	}
+
+	// Clone Message
+	cloned := *msg
+	if msg.Stats != nil {
+		stats := *msg.Stats
+		cloned.Stats = &stats
+		cloneMessageStatsPointers(msg.Stats, cloned.Stats)
+	}
+	return &cloned
+}
+
+func cloneMessageStatsPointers(src, dst *types.MessageStats) {
+	// Clone Pointer Fields
+	if src.EndTime != nil {
+		v := *src.EndTime
+		dst.EndTime = &v
+	}
+	if src.PromptTokens != nil {
+		v := *src.PromptTokens
+		dst.PromptTokens = &v
+	}
+	if src.GeneratedTokens != nil {
+		v := *src.GeneratedTokens
+		dst.GeneratedTokens = &v
+	}
+	if src.PromptPerSec != nil {
+		v := *src.PromptPerSec
+		dst.PromptPerSec = &v
+	}
+	if src.GeneratedPerSec != nil {
+		v := *src.GeneratedPerSec
+		dst.GeneratedPerSec = &v
+	}
+	if src.TimeToFirstToken != nil {
+		v := *src.TimeToFirstToken
+		dst.TimeToFirstToken = &v
+	}
+	if src.TimeToLastToken != nil {
+		v := *src.TimeToLastToken
+		dst.TimeToLastToken = &v
+	}
+}
--- a/backend/internal/api/handlers.go
+++ b/backend/internal/api/handlers.go
@@ -24,17 +24,19 @@ import (
 )

 type API struct {
-	logger  *logrus.Entry
-	store   store.Store
-	client  *client.Client
-	dataDir string
+	logger            *logrus.Entry
+	store             store.Store
+	client            *client.Client
+	dataDir           string
+	generationManager *generationManager
 }

 func New(s store.Store, dataDir string, logger *logrus.Logger) *API {
 	return &API{
-		store:   s,
-		dataDir: dataDir,
-		logger:  logger.WithField("service", "api"),
+		store:             s,
+		dataDir:           dataDir,
+		logger:            logger.WithField("service", "api"),
+		generationManager: newGenerationManager(),
 	}
 }

@@ -322,13 +324,18 @@ func (a *API) PostChat(w http.ResponseWriter, r *http.Request) {
 		return
 	}

-	// Send Message
-	responseStarted, err := a.sendMessage(r.Context(), w, chat.ID, genReq.Model, genReq.Prompt)
+	// Start Message
+	chunk, err := a.startMessageGeneration(chat.ID, genReq.Model, genReq.Prompt)
 	if err != nil {
-		log.WithError(err).WithField("chat_id", chat.ID).Error("failed to send message")
-		if !responseStarted {
-			http.Error(w, "Failed to send message", http.StatusInternalServerError)
-		}
+		log.WithError(err).WithField("chat_id", chat.ID).Error("failed to start message generation")
+		http.Error(w, "Failed to start message generation", http.StatusInternalServerError)
+		return
+	}
+
+	w.Header().Set("Content-Type", "application/json")
+	if err := json.NewEncoder(w).Encode(chunk); err != nil {
+		log.WithError(err).Error("failed to encode message generation response")
+		http.Error(w, "Failed to encode message generation response", http.StatusInternalServerError)
 	}
 }

@@ -395,6 +402,68 @@ func (a *API) GetChat(w http.ResponseWriter, r *http.Request) {
 	}
 }

+func (a *API) GetChatStream(w http.ResponseWriter, r *http.Request) {
+	log := a.logger.WithField("handler", "GetChatStreamHandler")
+
+	// Parse Chat ID
+	rawChatID := r.PathValue("chatId")
+	if rawChatID == "" {
+		log.Error("missing chat ID parameter")
+		http.Error(w, "Chat ID is required", http.StatusBadRequest)
+		return
+	}
+	chatID, err := uuid.Parse(rawChatID)
+	if err != nil {
+		log.WithError(err).Error("invalid chat ID format")
+		http.Error(w, "Invalid chat ID format", http.StatusBadRequest)
+		return
+	}
+
+	// Subscribe Before Snapshot
+	updates, unsubscribe, active := a.generationManager.subscribe(chatID)
+	defer unsubscribe()
+
+	// Get Chat Snapshot
+	chat, err := a.store.GetChat(chatID)
+	if err != nil {
+		log.WithError(err).WithField("chat_id", chatID).Error("failed to get chat")
+		http.Error(w, "Failed to get chat", http.StatusInternalServerError)
+		return
+	}
+
+	// Set Headers
+	w.Header().Set("Content-Type", "application/x-ndjson")
+	w.Header().Set("Cache-Control", "no-cache")
+	w.Header().Set("Connection", "keep-alive")
+	w.Header().Set("Transfer-Encoding", "chunked")
+	flushWriter := newFlushWriter(w)
+
+	// Send Snapshot
+	if err := json.NewEncoder(flushWriter).Encode(&MessageChunk{Chat: toChat(chat)}); err != nil {
+		log.WithError(err).WithField("chat_id", chatID).Warn("failed to send stream snapshot")
+		return
+	}
+	if !active {
+		return
+	}
+
+	// Forward Updates
+	for {
+		select {
+		case <-r.Context().Done():
+			return
+		case chunk, ok := <-updates:
+			if !ok {
+				return
+			}
+			if err := json.NewEncoder(flushWriter).Encode(chunk); err != nil {
+				log.WithError(err).WithField("chat_id", chatID).Warn("client stream disconnected")
+				return
+			}
+		}
+	}
+}
+
 func (a *API) PostChatMessage(w http.ResponseWriter, r *http.Request) {
 	log := a.logger.WithField("handler", "PostChatMessageHandler")

@@ -424,13 +493,22 @@ func (a *API) PostChatMessage(w http.ResponseWriter, r *http.Request) {
 		return
 	}

-	// Send Message
-	responseStarted, err := a.sendMessage(r.Context(), w, chatID, genReq.Model, genReq.Prompt)
+	// Start Message
+	chunk, err := a.startMessageGeneration(chatID, genReq.Model, genReq.Prompt)
 	if err != nil {
-		log.WithError(err).WithField("chat_id", chatID).Error("failed to send message")
-		if !responseStarted {
-			http.Error(w, "Failed to send message", http.StatusInternalServerError)
+		log.WithError(err).WithField("chat_id", chatID).Error("failed to start message generation")
+		if errors.Is(err, errGenerationActive) {
+			http.Error(w, "Chat generation already active", http.StatusConflict)
+		} else {
+			http.Error(w, "Failed to start message generation", http.StatusInternalServerError)
 		}
+		return
+	}
+
+	w.Header().Set("Content-Type", "application/json")
+	if err := json.NewEncoder(w).Encode(chunk); err != nil {
+		log.WithError(err).Error("failed to encode message generation response")
+		http.Error(w, "Failed to encode message generation response", http.StatusInternalServerError)
 	}
 }

@@ -456,99 +534,99 @@ func (a *API) getClient() (*client.Client, error) {
 	return a.client, nil
 }

-func (a *API) sendMessage(ctx context.Context, w http.ResponseWriter, chatID uuid.UUID, chatModel, userMessage string) (bool, error) {
+func (a *API) startMessageGeneration(chatID uuid.UUID, chatModel, userMessage string) (*MessageChunk, error) {
 	apiClient, err := a.getClient()
 	if err != nil {
-		return false, fmt.Errorf("failed to get client: %w", err)
+		return nil, fmt.Errorf("failed to get client: %w", err)
 	}

-	// Detach Request Context
-	ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), time.Minute*5)
-	defer cancel()
+	var chat *store.Chat
+	var userMsg *store.Message
+	var assistantMsg *store.Message
+	var initialChunk *MessageChunk

-	// Create User Message
-	userMsg := &store.Message{ChatID: chatID, Role: "user", Content: userMessage}
-	if err := a.store.SaveChatMessage(userMsg); err != nil {
-		return false, fmt.Errorf("failed to add user message to chat: %w", err)
-	}
+	// Start Generation - The manager reserves the chat before messages are
+	// persisted, preventing concurrent completions from creating duplicate rows.
+	if err := a.generationManager.start(chatID, func(_ *generation) error {
+		// Create User Message
+		userMsg = &store.Message{ChatID: chatID, Role: "user", Content: userMessage}
+		if err := a.store.SaveChatMessage(userMsg); err != nil {
+			return fmt.Errorf("failed to add user message to chat: %w", err)
+		}

-	// Get Chat History - Fetch before creating the in-progress assistant message so the
-	// LLM request does not include an empty assistant response prefill.
-	chat, err := a.store.GetChat(chatID)
-	if err != nil {
-		return false, fmt.Errorf("failed to get chat: %w", err)
-	}
+		// Get Chat History - Fetch before creating the in-progress assistant message so the
+		// LLM request does not include an empty assistant response prefill.
+		chat, err = a.store.GetChat(chatID)
+		if err != nil {
+			return fmt.Errorf("failed to get chat: %w", err)
+		}

-	// Add Assistant Response - TODO: Ensure InProgress Flag?
-	assistantMsg := &store.Message{ChatID: chatID, Role: "assistant"}
-	if err := a.store.SaveChatMessage(assistantMsg); err != nil {
-		return false, fmt.Errorf("failed to add assistant message to chat: %w", err)
-	}
+		// Add Assistant Response
+		assistantMsg = &store.Message{ChatID: chatID, Role: "assistant", Status: store.MessageStatusStreaming}
+		if err := a.store.SaveChatMessage(assistantMsg); err != nil {
+			return fmt.Errorf("failed to add assistant message to chat: %w", err)
+		}

-	// Set Headers
-	w.Header().Set("Content-Type", "application/x-ndjson")
-	w.Header().Set("Cache-Control", "no-cache")
-	w.Header().Set("Connection", "keep-alive")
-	w.Header().Set("Transfer-Encoding", "chunked")
-
-	// Create Flush Writer
-	flushWriter := newFlushWriter(w)
-
-	// Send Initial Chunk - User Message & Chat
-	if err := json.NewEncoder(flushWriter).Encode(&MessageChunk{
-		Chat:        toChatNoMessages(chat),
-		UserMessage: userMsg,
+		// Create Initial Chunk
+		initialChunk = &MessageChunk{
+			Chat:             toChatNoMessages(chat),
+			UserMessage:      userMsg,
+			AssistantMessage: assistantMsg,
+		}
+		return nil
+	}, func(gen *generation) {
+		a.runMessageGeneration(apiClient, chat, assistantMsg, chatModel, gen)
 	}); err != nil {
-		return false, fmt.Errorf("failed to send initial chunk: %w", err)
+		return nil, err
 	}
-	responseStarted := true
-	streamToClient := true
+
+	return initialChunk, nil
+}
+
+func (a *API) runMessageGeneration(apiClient *client.Client, chat *store.Chat, assistantMsg *store.Message, chatModel string, gen *generation) {
+	// Create Generation Context
+	ctx, cancel := context.WithTimeout(context.Background(), time.Minute*5)
+	defer cancel()

 	// Send Message
 	if _, err := apiClient.SendMessage(ctx, chat.Messages, chatModel, func(m *client.MessageChunk) error {
-		var apiMsgChunk MessageChunk
 		messageChanged := false

 		if m.Stats != nil {
 			messageChanged = true
 			assistantMsg.Stats = m.Stats
 		}
-
 		if m.Message != nil {
 			messageChanged = true
 			assistantMsg.Content += *m.Message
 		}
-
 		if m.Thinking != nil {
 			messageChanged = true
 			assistantMsg.Thinking += *m.Thinking
 		}

-		// Save Assistant Progress - Persist each streamed update so partial content
-		// survives client disconnects or upstream stream failures.
+		// Save And Broadcast Progress
 		if messageChanged {
 			if err := a.store.SaveChatMessage(assistantMsg); err != nil {
 				return fmt.Errorf("failed to save assistant progress: %w", err)
 			}
-			apiMsgChunk.AssistantMessage = assistantMsg
-		}
-
-		// Send Progress Chunk - If the browser disconnects, keep the detached
-		// generation running and continue saving streamed content to the store.
-		if streamToClient {
-			if err := json.NewEncoder(flushWriter).Encode(apiMsgChunk); err != nil {
-				streamToClient = false
-				a.logger.WithError(err).WithField("chat_id", chat.ID).Warn("client stream disconnected")
-			}
+			gen.broadcast(&MessageChunk{AssistantMessage: assistantMsg})
 		}

 		return nil
 	}); err != nil {
-		return responseStarted, fmt.Errorf("failed to generate text stream: %w", err)
+		assistantMsg.Status = store.MessageStatusFailed
+		if saveErr := a.store.SaveChatMessage(assistantMsg); saveErr != nil {
+			a.logger.WithError(saveErr).WithField("chat_id", chat.ID).Error("failed to save failed assistant message")
+		}
+		gen.broadcast(&MessageChunk{AssistantMessage: assistantMsg})
+		a.logger.WithError(err).WithField("chat_id", chat.ID).Error("failed to generate text stream")
+		return
 	}

 	// Summarize & Update Chat Title
 	if chat.Title == "" {
+		var err error
 		chat.Title, err = apiClient.CreateTitle(ctx, chat.Messages[0].Content, chatModel)
 		if err != nil {
 			a.logger.WithError(err).WithField("chat_id", chat.ID).Error("failed to create chat title")
@@ -557,20 +635,11 @@ func (a *API) sendMessage(ctx context.Context, w http.ResponseWriter, chatID uui
 		}
 	}

-	// Update Assistant Message
+	// Complete Assistant Message
+	assistantMsg.Status = store.MessageStatusComplete
 	if err := a.store.SaveChatMessage(assistantMsg); err != nil {
-		return responseStarted, fmt.Errorf("failed to save assistant message to chat: %w", err)
+		a.logger.WithError(err).WithField("chat_id", chat.ID).Error("failed to save assistant message")
+		return
 	}
-
-	// Send Final Chunk
-	if streamToClient {
-		if err := json.NewEncoder(flushWriter).Encode(&MessageChunk{
-			Chat:             toChatNoMessages(chat),
-			AssistantMessage: assistantMsg,
-		}); err != nil {
-			return responseStarted, fmt.Errorf("failed to send final chunk: %w", err)
-		}
-	}
-
-	return responseStarted, nil
+	gen.broadcast(&MessageChunk{Chat: toChatNoMessages(chat), AssistantMessage: assistantMsg})
 }
--- a/backend/internal/server/server.go
+++ b/backend/internal/server/server.go
@@ -48,6 +48,7 @@ func StartServer(settingsStore store.Store, dataDir, staticDir, listenAddress st
 	mux.HandleFunc("GET /api/chats", api.GetChats)
 	mux.HandleFunc("POST /api/chats", api.PostChat)
 	mux.HandleFunc("GET /api/chats/{chatId}", api.GetChat)
+	mux.HandleFunc("GET /api/chats/{chatId}/stream", api.GetChatStream)
 	mux.HandleFunc("POST /api/chats/{chatId}", api.PostChatMessage)
 	mux.HandleFunc("DELETE /api/chats/{chatId}", api.DeleteChat)

--- a/backend/internal/store/types.go
+++ b/backend/internal/store/types.go
@@ -28,6 +28,14 @@ type Chat struct {
 	Messages []*Message `json:"messages"`
 }

+type MessageStatus string
+
+const (
+	MessageStatusStreaming MessageStatus = "streaming"
+	MessageStatusComplete  MessageStatus = "complete"
+	MessageStatusFailed    MessageStatus = "failed"
+)
+
 type Message struct {
 	baseModel

@@ -35,5 +43,6 @@ type Message struct {
 	Role     string              `json:"role"`
 	Thinking string              `json:"thinking"`
 	Content  string              `json:"content"`
+	Status   MessageStatus       `json:"status,omitempty"`
 	Stats    *types.MessageStats `json:"stats,omitempty"`
 }