package stream import ( "context" "encoding/json" "fmt" "os" "sync" "time" "github.com/geo-platform/tenant-api/internal/shared/messaging/rabbitmq" ) 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 generationEventEnvelope struct { Source string `json:"source"` Event GenerationEvent `json:"event"` } type generationStream struct { snapshot GenerationEvent subscribers map[int]chan GenerationEvent nextSubscriber int completed bool } type GenerationHub struct { mu sync.Mutex streams map[int64]*generationStream rabbitMQ *rabbitmq.Client instanceID string consumerName string } func NewGenerationHub(rabbitMQClient *rabbitmq.Client) *GenerationHub { instanceID := fmt.Sprintf("%d-%d", os.Getpid(), time.Now().UnixNano()) return &GenerationHub{ streams: make(map[int64]*generationStream), rabbitMQ: rabbitMQClient, instanceID: instanceID, consumerName: "generation-stream-" + instanceID, } } func (h *GenerationHub) Run(ctx context.Context) { if h == nil || h.rabbitMQ == nil { return } go h.run(ctx) } func (h *GenerationHub) Seed(event GenerationEvent) { h.mu.Lock() defer h.mu.Unlock() state := h.ensureLocked(event.ArticleID) applySnapshotLocked(state, event) if state.completed && len(state.subscribers) == 0 { delete(h.streams, event.ArticleID) } } func (h *GenerationHub) StartEvent(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) Start(articleID, taskID int64, title string) { h.StartEvent(articleID, taskID, title) } 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 state.completed = false 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) run(ctx context.Context) { for { deliveries, ch, err := h.rabbitMQ.ConsumeGenerationEvents(h.consumerName) if err != nil { select { case <-ctx.Done(): return case <-time.After(2 * time.Second): continue } } closed := false for !closed { select { case <-ctx.Done(): _ = ch.Close() return case delivery, ok := <-deliveries: if !ok { closed = true continue } h.handleDelivery(delivery.Body) _ = delivery.Ack(false) } } _ = ch.Close() select { case <-ctx.Done(): return case <-time.After(2 * time.Second): } } } func (h *GenerationHub) handleDelivery(body []byte) { var envelope generationEventEnvelope if err := json.Unmarshal(body, &envelope); err != nil { return } if envelope.Source == h.instanceID { return } h.mu.Lock() state := h.streams[envelope.Event.ArticleID] if state == nil { h.mu.Unlock() return } applyRemoteEventLocked(state, envelope.Event) subscribers := collectSubscribersLocked(state) shouldDelete := state.completed && len(state.subscribers) == 0 if shouldDelete { delete(h.streams, envelope.Event.ArticleID) } h.mu.Unlock() for _, ch := range subscribers { select { case ch <- envelope.Event: default: } } } func (h *GenerationHub) publish(articleID int64, mutate func(*generationStream) GenerationEvent) { h.mu.Lock() state := h.ensureLocked(articleID) event := mutate(state) subscribers := collectSubscribersLocked(state) shouldDelete := state.completed && len(state.subscribers) == 0 if shouldDelete { delete(h.streams, articleID) } h.mu.Unlock() for _, ch := range subscribers { select { case ch <- event: default: } } if h.rabbitMQ != nil { payload, err := json.Marshal(generationEventEnvelope{ Source: h.instanceID, Event: event, }) if err == nil { _ = h.rabbitMQ.PublishGenerationEvent(context.Background(), payload) } } } 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 } func collectSubscribersLocked(state *generationStream) []chan GenerationEvent { subscribers := make([]chan GenerationEvent, 0, len(state.subscribers)) for _, ch := range state.subscribers { subscribers = append(subscribers, ch) } return subscribers } func applySnapshotLocked(state *generationStream, event GenerationEvent) { state.snapshot = GenerationEvent{ Type: "snapshot", ArticleID: event.ArticleID, TaskID: event.TaskID, Title: event.Title, Status: event.Status, Content: event.Content, Error: event.Error, Done: event.Done, UpdatedAt: event.UpdatedAt, } state.completed = event.Done } func applyRemoteEventLocked(state *generationStream, event GenerationEvent) { switch event.Type { case "delta": state.completed = false state.snapshot.Type = "snapshot" state.snapshot.ArticleID = event.ArticleID state.snapshot.TaskID = event.TaskID state.snapshot.Title = event.Title state.snapshot.Status = event.Status if event.Content != "" { state.snapshot.Content = event.Content } else { state.snapshot.Content += event.Delta } state.snapshot.UpdatedAt = event.UpdatedAt case "completed": applySnapshotLocked(state, GenerationEvent{ Type: "snapshot", ArticleID: event.ArticleID, TaskID: event.TaskID, Title: event.Title, Status: "completed", Content: event.Content, Done: true, UpdatedAt: event.UpdatedAt, }) case "error": applySnapshotLocked(state, GenerationEvent{ Type: "snapshot", ArticleID: event.ArticleID, TaskID: event.TaskID, Title: event.Title, Status: "failed", Error: event.Error, Done: true, UpdatedAt: event.UpdatedAt, }) default: applySnapshotLocked(state, GenerationEvent{ Type: "snapshot", ArticleID: event.ArticleID, TaskID: event.TaskID, Title: event.Title, Status: event.Status, Content: event.Content, Error: event.Error, Done: event.Done, UpdatedAt: event.UpdatedAt, }) } }