geo/server/internal/shared/stream/generation.go

package stream

import (
	"sync"
	"time"
)

type GenerationEvent struct {
	Type      string    `json:"type"`
	ArticleID int64     `json:"article_id"`
	TaskID    int64     `json:"task_id"`
	Title     string    `json:"title,omitempty"`
	Status    string    `json:"status,omitempty"`
	Delta     string    `json:"delta,omitempty"`
	Content   string    `json:"content,omitempty"`
	Error     string    `json:"error,omitempty"`
	Done      bool      `json:"done,omitempty"`
	UpdatedAt time.Time `json:"updated_at"`
}

type generationStream struct {
	snapshot       GenerationEvent
	subscribers    map[int]chan GenerationEvent
	nextSubscriber int
	completed      bool
}

type GenerationHub struct {
	mu      sync.Mutex
	streams map[int64]*generationStream
}

func NewGenerationHub() *GenerationHub {
	return &GenerationHub{
		streams: make(map[int64]*generationStream),
	}
}

func (h *GenerationHub) Start(articleID, taskID int64, title string) {
	h.publish(articleID, func(state *generationStream) GenerationEvent {
		now := time.Now()
		state.completed = false
		state.snapshot = GenerationEvent{
			Type:      "snapshot",
			ArticleID: articleID,
			TaskID:    taskID,
			Title:     title,
			Status:    "generating",
			UpdatedAt: now,
		}
		return GenerationEvent{
			Type:      "status",
			ArticleID: articleID,
			TaskID:    taskID,
			Title:     title,
			Status:    "generating",
			UpdatedAt: now,
		}
	})
}

func (h *GenerationHub) AppendDelta(articleID, taskID int64, title, delta string) {
	h.publish(articleID, func(state *generationStream) GenerationEvent {
		now := time.Now()
		state.snapshot.Type = "snapshot"
		state.snapshot.ArticleID = articleID
		state.snapshot.TaskID = taskID
		state.snapshot.Title = title
		state.snapshot.Status = "generating"
		state.snapshot.Content += delta
		state.snapshot.UpdatedAt = now

		return GenerationEvent{
			Type:      "delta",
			ArticleID: articleID,
			TaskID:    taskID,
			Title:     title,
			Status:    "generating",
			Delta:     delta,
			Content:   state.snapshot.Content,
			UpdatedAt: now,
		}
	})
}

func (h *GenerationHub) Complete(articleID, taskID int64, title, content string) {
	h.publish(articleID, func(state *generationStream) GenerationEvent {
		now := time.Now()
		state.completed = true
		state.snapshot = GenerationEvent{
			Type:      "snapshot",
			ArticleID: articleID,
			TaskID:    taskID,
			Title:     title,
			Status:    "completed",
			Content:   content,
			Done:      true,
			UpdatedAt: now,
		}

		return GenerationEvent{
			Type:      "completed",
			ArticleID: articleID,
			TaskID:    taskID,
			Title:     title,
			Status:    "completed",
			Content:   content,
			Done:      true,
			UpdatedAt: now,
		}
	})
}

func (h *GenerationHub) Fail(articleID, taskID int64, title, errMsg string) {
	h.publish(articleID, func(state *generationStream) GenerationEvent {
		now := time.Now()
		state.completed = true
		state.snapshot.Type = "snapshot"
		state.snapshot.ArticleID = articleID
		state.snapshot.TaskID = taskID
		state.snapshot.Title = title
		state.snapshot.Status = "failed"
		state.snapshot.Error = errMsg
		state.snapshot.Done = true
		state.snapshot.UpdatedAt = now

		return GenerationEvent{
			Type:      "error",
			ArticleID: articleID,
			TaskID:    taskID,
			Title:     title,
			Status:    "failed",
			Error:     errMsg,
			Done:      true,
			UpdatedAt: now,
		}
	})
}

func (h *GenerationHub) Subscribe(articleID int64) (<-chan GenerationEvent, *GenerationEvent, func()) {
	h.mu.Lock()
	state := h.streams[articleID]
	if state == nil {
		h.mu.Unlock()
		return nil, nil, func() {}
	}

	id := state.nextSubscriber
	state.nextSubscriber++

	ch := make(chan GenerationEvent, 64)
	state.subscribers[id] = ch
	snapshot := state.snapshot
	h.mu.Unlock()

	cancel := func() {
		h.mu.Lock()
		defer h.mu.Unlock()

		current := h.streams[articleID]
		if current == nil {
			close(ch)
			return
		}

		delete(current.subscribers, id)
		shouldDelete := current.completed && len(current.subscribers) == 0
		if shouldDelete {
			delete(h.streams, articleID)
		}

		close(ch)
	}

	return ch, &snapshot, cancel
}

func (h *GenerationHub) Snapshot(articleID int64) (*GenerationEvent, bool) {
	h.mu.Lock()
	defer h.mu.Unlock()

	state := h.streams[articleID]
	if state == nil {
		return nil, false
	}

	snapshot := state.snapshot
	return &snapshot, true
}

func (h *GenerationHub) publish(articleID int64, mutate func(*generationStream) GenerationEvent) {
	h.mu.Lock()
	state := h.ensureLocked(articleID)
	event := mutate(state)

	subscribers := make([]chan GenerationEvent, 0, len(state.subscribers))
	for _, ch := range state.subscribers {
		subscribers = append(subscribers, ch)
	}
	h.mu.Unlock()

	for _, ch := range subscribers {
		select {
		case ch <- event:
		default:
		}
	}
}

func (h *GenerationHub) ensureLocked(articleID int64) *generationStream {
	state := h.streams[articleID]
	if state != nil {
		return state
	}

	state = &generationStream{
		subscribers: make(map[int]chan GenerationEvent),
	}
	h.streams[articleID] = state
	return state
}