moteva/server/internal/logic/agent_websocket_logic.go

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

package logic

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
	"time"

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

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

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

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

func (l *AgentWebsocketLogic) AgentWebsocket(req *types.ProjectIdRequest) error {
	return nil
}

type socketEnvelope struct {
	Role          string          `json:"role"`
	Type          string          `json:"type"`
	ReqUUID       string          `json:"req_uuid"`
	SendTimestamp int64           `json:"send_timestamp"`
	Data          socketStartData `json:"data"`
}

type socketStartData struct {
	Version               int64               `json:"version"`
	Event                 string              `json:"event"`
	ThreadID              string              `json:"thread_id"`
	ProjectID             string              `json:"project_id"`
	Messages              []socketChatMessage `json:"messages"`
	ToolConfig            socketToolConfig    `json:"tool_config"`
	ThreadIDType          int64               `json:"thread_id_type"`
	EnableMultimodalInput int64               `json:"enable_multimodal_input"`
	EnableWebSearch       bool                `json:"enable_web_search"`
	ImageModel            string              `json:"image_model"`
	ImageSize             string              `json:"image_size"`
	Mode                  string              `json:"mode"`
}

type socketChatMessage struct {
	Role     string          `json:"role"`
	Contents []socketContent `json:"contents"`
}

type socketContent struct {
	Type        string `json:"type"`
	Text        string `json:"text"`
	TextSource  string `json:"text_source"`
	ImageURL    string `json:"image_url"`
	ImageWidth  int64  `json:"image_width"`
	ImageHeight int64  `json:"image_height"`
}

type socketToolConfig struct {
	PreferToolCategories map[string][]string `json:"prefer_tool_categories"`
}

type socketMessage struct {
	Type           string `json:"type"`
	MsgID          string `json:"msg_id,omitempty"`
	ThreadID       string `json:"thread_id,omitempty"`
	ActionID       string `json:"action_id,omitempty"`
	StepID         string `json:"step_id,omitempty"`
	ParentThreadID string `json:"parent_thread_id,omitempty"`
	ParentActionID string `json:"parent_action_id,omitempty"`
	AgentName      string `json:"agent_name,omitempty"`
	Data           any    `json:"data,omitempty"`
	ContentBlock   any    `json:"content_block,omitempty"`
	Project        any    `json:"project,omitempty"`
	Text           string `json:"text,omitempty"`
	Name           string `json:"name,omitempty"`
	ToolHint       string `json:"tool_hint,omitempty"`
	Status         string `json:"status,omitempty"`
}

var agentSocketUpgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool { return true },
}

func (l *AgentWebsocketLogic) Stream(w http.ResponseWriter, r *http.Request, req *types.ProjectIdRequest) error {
	conn, err := agentSocketUpgrader.Upgrade(w, r, nil)
	if err != nil {
		return err
	}
	defer conn.Close()

	threadID := strings.TrimSpace(r.URL.Query().Get("threadId"))
	if threadID == "" {
		threadID = strings.TrimSpace(r.URL.Query().Get("thread_id"))
	}

	if threadID == "" {
		thread, err := l.readStartAndRun(conn, req.Id)
		if err != nil {
			return err
		}
		threadID = thread.ThreadID
	}

	ackDone := make(chan struct{})
	go l.readAcknowledgements(conn, ackDone)
	defer close(ackDone)

	return l.streamThread(conn, req.Id, threadID)
}

func (l *AgentWebsocketLogic) readStartAndRun(conn *websocket.Conn, projectID string) (design.AgentThread, error) {
	_, payload, err := conn.ReadMessage()
	if err != nil {
		return design.AgentThread{}, err
	}

	var envelope socketEnvelope
	if err := json.Unmarshal(payload, &envelope); err != nil {
		return design.AgentThread{}, err
	}
	if envelope.Data.Event != "start" {
		return design.AgentThread{}, fmt.Errorf("first agent socket message must be a start event")
	}

	req := envelope.Data.toDomainChat(projectID)
	return l.svcCtx.DesignService.AgentChat(l.ctx, projectID, req)
}

func (data socketStartData) toDomainChat(projectID string) design.AgentChatRequest {
	messages := make([]design.AgentChatMessage, 0, len(data.Messages))
	for _, message := range data.Messages {
		contents := make([]design.AgentContent, 0, len(message.Contents))
		for _, content := range message.Contents {
			contents = append(contents, design.AgentContent{
				Type:        content.Type,
				Text:        content.Text,
				TextSource:  content.TextSource,
				ImageURL:    content.ImageURL,
				ImageWidth:  content.ImageWidth,
				ImageHeight: content.ImageHeight,
			})
		}
		messages = append(messages, design.AgentChatMessage{
			Role:     message.Role,
			Contents: contents,
		})
	}
	return design.AgentChatRequest{
		Messages: messages,
		ToolConfig: design.AgentToolConfig{
			PreferToolCategories: data.ToolConfig.PreferToolCategories,
		},
		ThreadID:              data.ThreadID,
		ThreadIDType:          data.ThreadIDType,
		EnableMultimodalInput: data.EnableMultimodalInput,
		EnableWebSearch:       data.EnableWebSearch || hasSocketPreferredToolCategory(data.ToolConfig.PreferToolCategories, "SEARCH"),
		ImageModel:            data.ImageModel,
		ImageSize:             data.ImageSize,
		Mode:                  data.Mode,
	}
}

func hasSocketPreferredToolCategory(categories map[string][]string, key string) bool {
	for category, items := range categories {
		if strings.EqualFold(category, key) && len(items) > 0 {
			return true
		}
	}
	return false
}

func (l *AgentWebsocketLogic) readAcknowledgements(conn *websocket.Conn, done <-chan struct{}) {
	for {
		select {
		case <-done:
			return
		default:
			var envelope socketEnvelope
			if err := conn.ReadJSON(&envelope); err != nil {
				return
			}
		}
	}
}

func (l *AgentWebsocketLogic) streamThread(conn *websocket.Conn, projectID string, threadID string) error {
	ticker := time.NewTicker(600 * time.Millisecond)
	defer ticker.Stop()

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

	actionID := threadID
	lastFingerprint := ""
	messageEvents := newMessageEventState()

	if err := writeSocketJSON(conn, socketMessage{
		Type:      "thread_start",
		MsgID:     newSocketID(),
		ThreadID:  threadID,
		AgentName: "pi-agent",
	}); err != nil {
		return err
	}
	if err := writeSocketJSON(conn, socketMessage{
		Type:     "action_start",
		MsgID:    newSocketID(),
		ThreadID: threadID,
		ActionID: actionID,
	}); err != nil {
		return err
	}

	for {
		project, err := l.svcCtx.DesignService.GetProject(l.ctx, projectID)
		if err != nil {
			return err
		}

		apiProject := toAPIProject(project)
		fingerprint := projectEventFingerprint(project)
		if fingerprint != lastFingerprint {
			if err := writeSocketJSON(conn, socketMessage{
				Type:    "project",
				MsgID:   newSocketID(),
				Project: ProjectEventPayload{Project: apiProject},
			}); err != nil {
				return err
			}
			lastFingerprint = fingerprint
		}

		hiddenThreadIds := hiddenAgentThreadIDs(project.Messages)
		for _, message := range project.Messages {
			if threadID != "" && message.ThreadID != "" && message.ThreadID != threadID {
				continue
			}
			if !messageEvents.shouldEmit(message, time.Now()) {
				continue
			}
			if message.Type != "text_extraction_result" && hiddenThreadIds[message.ThreadID] {
				continue
			}
			if err := l.streamMessage(conn, threadID, actionID, message); err != nil {
				return err
			}
		}

		threadStarted, threadFinished, threadFailed := projectThreadState(project.Messages, threadID)
		if (threadID != "" && threadStarted && (threadFinished || threadFailed)) || (threadID == "" && (project.Status == design.StatusReady || project.Status == design.StatusFailed)) {
			status := string(project.Status)
			if threadID != "" && threadStarted {
				status = "ready"
				if threadFailed {
					status = "failed"
				}
			}
			if err := writeSocketJSON(conn, socketMessage{
				Type:    "done",
				MsgID:   newSocketID(),
				Project: ProjectEventPayload{Project: apiProject},
			}); err != nil {
				return err
			}
			return writeSocketJSON(conn, socketMessage{
				Type:     "thread_end",
				MsgID:    newSocketID(),
				ThreadID: threadID,
				ActionID: actionID,
				Status:   status,
			})
		}

		select {
		case <-l.ctx.Done():
			return l.ctx.Err()
		case <-deadline.C:
			return writeSocketJSON(conn, socketMessage{
				Type:    "timeout",
				MsgID:   newSocketID(),
				Project: ProjectEventPayload{Project: apiProject},
			})
		case <-ticker.C:
		}
	}
}

func textExtractionThreadState(messages []design.Message, threadID string) (started bool, finished bool, failed bool) {
	for _, message := range messages {
		if message.ThreadID != threadID {
			continue
		}
		if message.Name == "text_extraction" || message.Type == "text_extraction_result" {
			started = true
		}
		if message.Type == "text_extraction_result" {
			finished = true
		}
		if message.Role == "error" && (message.Name == "text_extraction" || message.Title == "文字提取失败" || message.Title == "文字提取超时") {
			failed = true
		}
	}
	return started, finished, failed
}

func (l *AgentWebsocketLogic) streamMessage(conn *websocket.Conn, threadID string, actionID string, message design.Message) error {
	if message.Type == "text_extraction_result" {
		return writeSocketJSON(conn, socketMessage{
			Type:     "text_extraction",
			MsgID:    newSocketID(),
			ThreadID: threadID,
			ActionID: actionID,
			StepID:   fallbackStepID(message),
			Data:     parseJSONOrString(message.Content),
			Status:   coalesce(message.Status, "success"),
		})
	}
	if isHiddenAgentMessage(message) {
		return nil
	}

	if message.Role != "user" {
		if err := writeSocketJSON(conn, socketMessage{
			Type:  "thinking",
			MsgID: newSocketID(),
			Data:  toAPIMessage(message),
		}); err != nil {
			return err
		}
	}

	switch message.Role {
	case "user":
		return writeSocketJSON(conn, socketMessage{
			Type:     "user",
			MsgID:    newSocketID(),
			ThreadID: threadID,
			ActionID: actionID,
			StepID:   fallbackStepID(message),
			Text:     message.Content,
			Data: map[string]any{
				"input": message.Content,
			},
		})
	case "tool":
		if message.Type == "tool_use" || message.Name == "generate_media" {
			return writeSocketJSON(conn, socketMessage{
				Type:     "content_block_start",
				MsgID:    newSocketID(),
				ThreadID: threadID,
				ActionID: actionID,
				StepID:   fallbackStepID(message),
				ContentBlock: map[string]any{
					"id":      fallbackStepID(message),
					"type":    "tool_use",
					"name":    coalesce(message.Name, message.Title),
					"content": parseJSONOrString(message.Content),
					"status":  coalesce(message.Status, "success"),
				},
				Name:     coalesce(message.Name, message.Title),
				ToolHint: coalesce(message.ToolHint, "generate_image_gpt_image_2"),
				Status:   coalesce(message.Status, "success"),
			})
		}
		return writeSocketJSON(conn, socketMessage{
			Type:     "step_start",
			MsgID:    newSocketID(),
			ThreadID: threadID,
			ActionID: actionID,
			StepID:   fallbackStepID(message),
			Data: map[string]any{
				"title":   message.Title,
				"content": parseJSONOrString(message.Content),
			},
		})
	case "assistant":
		if strings.TrimSpace(message.Content) == "" && strings.TrimSpace(message.ToolHint) != "" {
			return writeSocketJSON(conn, socketMessage{
				Type:     "step_start",
				MsgID:    newSocketID(),
				ThreadID: threadID,
				ActionID: actionID,
				StepID:   fallbackStepID(message),
				ToolHint: message.ToolHint,
				Data: map[string]any{
					"tool_hint": message.ToolHint,
				},
			})
		}
		return writeSocketJSON(conn, socketMessage{
			Type:     "content_block_start",
			MsgID:    newSocketID(),
			ThreadID: threadID,
			ActionID: actionID,
			StepID:   fallbackStepID(message),
			ContentBlock: map[string]any{
				"id":   fallbackStepID(message),
				"type": "text",
				"text": message.Content,
			},
			Text: message.Content,
		})
	case "error":
		return writeSocketJSON(conn, socketMessage{
			Type:     "content_block_start",
			MsgID:    newSocketID(),
			ThreadID: threadID,
			ActionID: actionID,
			StepID:   fallbackStepID(message),
			ContentBlock: map[string]any{
				"id":         fallbackStepID(message),
				"type":       "error",
				"error_text": message.Content,
			},
			Status: "failed",
		})
	default:
		return nil
	}
}

func writeSocketJSON(conn *websocket.Conn, payload any) error {
	return conn.WriteJSON(payload)
}

func newSocketID() string {
	return fmt.Sprintf("%d", time.Now().UnixNano())
}

func fallbackStepID(message design.Message) string {
	if strings.TrimSpace(message.StepID) != "" {
		return message.StepID
	}
	return message.ID
}

func coalesce(values ...string) string {
	for _, value := range values {
		if strings.TrimSpace(value) != "" {
			return value
		}
	}
	return ""
}

func parseJSONOrString(value string) any {
	value = strings.TrimSpace(value)
	if value == "" {
		return ""
	}
	var payload any
	if err := json.Unmarshal([]byte(value), &payload); err != nil {
		return value
	}
	return payload
}