moteva/server/internal/logic/project_events_logic.go

// Code scaffolded by goctl. Safe to edit.
// goctl 1.10.1

package logic

import (
	"context"
	"fmt"
	"strings"
	"time"

	"img_infinite_canvas/internal/domain/design"
	"img_infinite_canvas/internal/svc"
	"img_infinite_canvas/internal/types"

	"github.com/zeromicro/go-zero/core/logx"
)

type ProjectEventsLogic struct {
	logx.Logger
	ctx    context.Context
	svcCtx *svc.ServiceContext
}

func NewProjectEventsLogic(ctx context.Context, svcCtx *svc.ServiceContext) *ProjectEventsLogic {
	return &ProjectEventsLogic{
		Logger: logx.WithContext(ctx),
		ctx:    ctx,
		svcCtx: svcCtx,
	}
}

func (l *ProjectEventsLogic) ProjectEvents(req *types.ProjectIdRequest) error {
	return l.Stream(req, "", func(string, any) error {
		return nil
	})
}

type ProjectEventEmitter func(event string, payload any) error

type ProjectEventPayload struct {
	Project types.Project `json:"project"`
}

func (l *ProjectEventsLogic) Stream(req *types.ProjectIdRequest, threadID string, emit ProjectEventEmitter) error {
	fallbackPoll := 10 * time.Second
	if l.svcCtx != nil {
		fallbackPoll = l.svcCtx.RealtimeFallbackPoll()
	}
	ticker := time.NewTicker(fallbackPoll)
	defer ticker.Stop()

	deadline := time.NewTimer(20 * time.Minute)
	defer deadline.Stop()

	lastFingerprint := ""
	messageEvents := newMessageEventState()
	threadID = strings.TrimSpace(threadID)
	var apiProject types.Project

	emitCurrent := func() (bool, error) {
		project, err := l.svcCtx.DesignService.GetProject(l.ctx, req.Id)
		if err != nil {
			return false, err
		}

		apiProject = toAPIProject(project)
		fingerprint := projectEventFingerprint(project)
		if fingerprint != lastFingerprint {
			if err := emit("project", ProjectEventPayload{Project: apiProject}); err != nil {
				return false, err
			}
			lastFingerprint = fingerprint
		}

		hiddenThreadIds := hiddenAgentThreadIDs(project.Messages)
		for _, message := range project.Messages {
			if !shouldEmitThreadMessage(message, threadID) {
				continue
			}
			if !messageEvents.shouldEmit(message, time.Now()) {
				continue
			}
			if message.Type == "text_extraction_result" {
				if err := emit("text_extraction", parseJSONOrString(message.Content)); err != nil {
					return false, err
				}
				continue
			}
			if hiddenThreadIds[message.ThreadID] {
				continue
			}
			if message.Role != "user" {
				apiMessage := toAPIMessage(message)
				if err := emit("thinking", apiMessage); err != nil {
					return false, err
				}
			}
		}

		threadStarted, threadFinished, threadFailed := projectThreadState(project.Messages, threadID)
		threadHasUnfinishedWork := hasUnfinishedThreadWork(project.Messages, threadID)
		projectSettled := !hasActiveCanvasWorkNode(project.Nodes) && (project.Status == "ready" || project.Status == "failed")
		if (threadID != "" && threadStarted && (threadFinished || threadFailed || (projectSettled && !threadHasUnfinishedWork))) || (threadID == "" && projectSettled) {
			return true, nil
		}
		return false, nil
	}

	done, err := emitCurrent()
	if err != nil {
		return err
	}
	if done {
		return emit("done", ProjectEventPayload{Project: apiProject})
	}

	var events <-chan any
	var closeSubscription func() error
	if l.svcCtx != nil && l.svcCtx.RealtimeBus != nil {
		sub, err := l.svcCtx.RealtimeBus.Subscribe(l.ctx, req.Id)
		if err != nil {
			l.Errorf("subscribe project events failed: %v", err)
		} else if sub != nil {
			eventCh := make(chan any, 32)
			events = eventCh
			closeSubscription = sub.Close
			go func() {
				defer close(eventCh)
				for event := range sub.C() {
					select {
					case eventCh <- event:
					case <-l.ctx.Done():
						return
					}
				}
			}()
			defer closeSubscription()
		}
	}

	for {
		select {
		case <-l.ctx.Done():
			return l.ctx.Err()
		case <-deadline.C:
			return emit("timeout", ProjectEventPayload{Project: apiProject})
		case _, ok := <-events:
			if !ok {
				events = nil
				continue
			}
			done, err := emitCurrent()
			if err != nil {
				return err
			}
			if done {
				return emit("done", ProjectEventPayload{Project: apiProject})
			}
		case <-ticker.C:
			done, err := emitCurrent()
			if err != nil {
				return err
			}
			if done {
				return emit("done", ProjectEventPayload{Project: apiProject})
			}
		}
	}
}

func shouldEmitThreadMessage(message design.Message, threadID string) bool {
	if threadID == "" {
		return true
	}
	if message.ThreadID == "" {
		return false
	}
	return message.ThreadID == threadID
}

func messageEventFingerprint(message design.Message) string {
	return fmt.Sprintf("%s:%s:%s:%s:%s:%s:%s:%s:%s", message.Role, message.Type, message.Title, message.Content, message.Name, message.ToolHint, message.Status, message.ThreadID, message.CreatedAt)
}

func hasActiveCanvasWorkNode(nodes []design.Node) bool {
	for _, node := range nodes {
		switch node.Status {
		case "generating", "text-extracting", "uploading", "mockup-modeling":
			return true
		}
	}
	return false
}

func projectThreadState(messages []design.Message, threadID string) (started bool, finished bool, failed bool) {
	threadID = strings.TrimSpace(threadID)
	if threadID == "" {
		return false, false, false
	}
	startIndex := latestThreadTaskStartIndex(messages, threadID)
	if startIndex < 0 {
		return false, false, false
	}
	for index := startIndex; index < len(messages); index++ {
		message := messages[index]
		if strings.TrimSpace(message.ThreadID) != threadID {
			continue
		}
		started = true
		status := strings.ToLower(strings.TrimSpace(message.Status))
		if status == "cancelled" || message.Type == "text_extraction_result" || isFinalToolThreadMessage(message, status) || isFinalAssistantThreadMessage(message, status) {
			finished = true
		}
		if status == "failed" || message.Role == "error" {
			failed = true
		}
	}
	if hasUnfinishedThreadWorkFrom(messages, threadID, startIndex) {
		finished = false
	}
	return started, finished, failed
}

func hasUnfinishedThreadWork(messages []design.Message, threadID string) bool {
	threadID = strings.TrimSpace(threadID)
	if threadID == "" {
		return false
	}
	return hasUnfinishedThreadWorkFrom(messages, threadID, latestThreadTaskStartIndex(messages, threadID))
}

func hasUnfinishedThreadWorkFrom(messages []design.Message, threadID string, startIndex int) bool {
	if startIndex < 0 {
		return false
	}
	running := make(map[string]bool)
	for index := startIndex; index < len(messages); index++ {
		message := messages[index]
		if strings.TrimSpace(message.ThreadID) != threadID {
			continue
		}
		key := threadWorkKey(message)
		if key == "" {
			continue
		}
		status := strings.ToLower(strings.TrimSpace(message.Status))
		switch status {
		case "running", "streaming", "submitted":
			running[key] = true
		case "success", "failed", "cancelled":
			delete(running, key)
		}
	}
	return len(running) > 0
}

func latestThreadTaskStartIndex(messages []design.Message, threadID string) int {
	threadID = strings.TrimSpace(threadID)
	if threadID == "" {
		return -1
	}
	firstIndex := -1
	latestUserIndex := -1
	for index, message := range messages {
		if strings.TrimSpace(message.ThreadID) != threadID {
			continue
		}
		if firstIndex < 0 {
			firstIndex = index
		}
		if strings.EqualFold(strings.TrimSpace(message.Role), "user") {
			latestUserIndex = index
		}
	}
	if latestUserIndex >= 0 {
		return latestUserIndex
	}
	return firstIndex
}

func threadWorkKey(message design.Message) string {
	role := strings.ToLower(strings.TrimSpace(message.Role))
	if role == "" {
		role = strings.ToLower(strings.TrimSpace(message.Type))
	}
	if role == "user" {
		return ""
	}
	if name := strings.TrimSpace(message.Name); name != "" {
		return role + ":name:" + name
	}
	if hint := strings.TrimSpace(message.ToolHint); hint != "" {
		if prefix, _, ok := strings.Cut(hint, ":"); ok {
			hint = prefix
		}
		return role + ":hint:" + hint
	}
	if title := strings.TrimSpace(message.Title); title != "" {
		return role + ":title:" + title
	}
	return ""
}

func isFinalToolThreadMessage(message design.Message, status string) bool {
	if status != "success" || message.Role != "tool" {
		return false
	}
	return message.Name == "mockup_model" || message.ToolHint == "mockup_model"
}

func isFinalAssistantThreadMessage(message design.Message, status string) bool {
	if status != "success" || message.Role != "assistant" {
		return false
	}
	if message.Name == "agent_planner" || strings.HasPrefix(message.ToolHint, "agent_planner") {
		return false
	}
	return true
}