Agent-System/internal/server/sse.go

package server

import (
	"encoding/json"
	"fmt"
	"net/http"
	"strconv"
	"sync"
	"time"

	"chat-agent-system/internal/bizdev"
	"chat-agent-system/internal/stream"
)

// StreamChatRequest SSE流式请求参数
type StreamChatRequest struct {
	SessionID      string `json:"session_id"`
	Message        string `json:"message"`
	IncludeDetails bool   `json:"include_details"` // 是否推送底层长篇推演草稿 (按需展开，默认 false)
	ProjectName    string `json:"project_name,omitempty"`
	CustomerName   string `json:"customer_name,omitempty"`
	DeviceType     string `json:"device_type,omitempty"`
	IntendedUse    string `json:"intended_use,omitempty"`
	TargetMarket   string `json:"target_market,omitempty"`
	WorkDir        string `json:"work_dir,omitempty"`
	UserName       string `json:"user_name,omitempty"`
	UserRole       string `json:"user_role,omitempty"`
	UserID         int64  `json:"user_id,omitempty"`
	ProjectID      int64  `json:"project_id,omitempty"`
}

func (s *Server) handleChatStream(w http.ResponseWriter, r *http.Request) {
	flusher, ok := w.(http.Flusher)
	if !ok {
		http.Error(w, "streaming unsupported by underlying connection", http.StatusInternalServerError)
		return
	}

	var req StreamChatRequest
	if r.Method == http.MethodPost {
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			http.Error(w, "invalid request body: "+err.Error(), http.StatusBadRequest)
			return
		}
	} else {
		// 允许 GET query params 方便 EventSource 及浏览器直接调试
		q := r.URL.Query()
		req.SessionID = q.Get("session_id")
		req.Message = q.Get("message")
		req.ProjectName = q.Get("project_name")
		req.CustomerName = q.Get("customer_name")
		req.DeviceType = q.Get("device_type")
		req.IntendedUse = q.Get("intended_use")
		req.TargetMarket = q.Get("target_market")
		req.WorkDir = q.Get("work_dir")
		req.UserName = q.Get("user_name")
		req.UserRole = q.Get("user_role")
		if inc := q.Get("include_details"); inc != "" {
			req.IncludeDetails, _ = strconv.ParseBool(inc)
		}
	}

	if req.Message == "" {
		http.Error(w, "message cannot be empty", http.StatusBadRequest)
		return
	}

	// 1. 设置 SSE 响应头
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	w.Header().Set("Access-Control-Allow-Origin", "*")
	w.Header().Set("X-Accel-Buffering", "no") // 禁用 Nginx 等反代缓冲
	flusher.Flush()

	// 2. 准备会话
	sess := s.getOrCreateSession(
		req.SessionID,
		req.ProjectName,
		req.CustomerName,
		req.DeviceType,
		req.IntendedUse,
		req.TargetMarket,
		req.WorkDir,
		req.UserName,
		req.UserRole,
	)

	// 3. 构建线程安全的 SSE 监听分发闭包
	var writeMu sync.Mutex
	listener := func(ev stream.StreamEvent) {
		writeMu.Lock()
		defer writeMu.Unlock()

		data, err := json.Marshal(ev)
		if err != nil {
			return
		}
		// 按照标准 SSE 协议规范发射: event: <type>\ndata: <json>\n\n
		fmt.Fprintf(w, "event: %s\ndata: %s\n\n", ev.Type, string(data))
		flusher.Flush()
	}

	// 4. 发起流式执行
	_, err := sess.SendMessageStream(r.Context(), req.Message, listener, req.IncludeDetails)
	if err != nil {
		writeMu.Lock()
		errEv := stream.StreamEvent{
			Type:      stream.EventError,
			Timestamp: time.Now(),
			Error:     err.Error(),
		}
		data, _ := json.Marshal(errEv)
		fmt.Fprintf(w, "event: %s\ndata: %s\n\n", errEv.Type, string(data))
		flusher.Flush()
		writeMu.Unlock()
	}
}

func (s *Server) handleRunTaskStream(w http.ResponseWriter, r *http.Request) {
	flusher, ok := w.(http.Flusher)
	if !ok {
		http.Error(w, "streaming unsupported", http.StatusInternalServerError)
		return
	}

	var req TaskRunRequestDTO
	includeDetails := false
	if r.Method == http.MethodPost {
		_ = json.NewDecoder(r.Body).Decode(&req)
	} else {
		q := r.URL.Query()
		req.TaskKey = q.Get("task_key")
		req.Stage = q.Get("stage")
		req.ProjectName = q.Get("project_name")
		req.CustomerName = q.Get("customer_name")
		req.DeviceType = q.Get("device_type")
		req.IntendedUse = q.Get("intended_use")
		req.TargetMarket = q.Get("target_market")
		req.Backend = q.Get("backend")
		req.AutoRepair, _ = strconv.ParseBool(q.Get("auto_repair"))
		if inc := q.Get("include_details"); inc != "" {
			includeDetails, _ = strconv.ParseBool(inc)
		}
	}

	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	w.Header().Set("Access-Control-Allow-Origin", "*")
	flusher.Flush()

	var writeMu sync.Mutex
	listener := func(ev stream.StreamEvent) {
		writeMu.Lock()
		defer writeMu.Unlock()
		data, err := json.Marshal(ev)
		if err != nil {
			return
		}
		fmt.Fprintf(w, "event: %s\ndata: %s\n\n", ev.Type, string(data))
		flusher.Flush()
	}

	ctxWithS, streamer := stream.WithListener(r.Context(), includeDetails, listener)
	start := time.Now()

	if req.Stage != "" {
		_, err := s.svc.ExecuteStage(ctxWithS, bizdev.StageRunRequest{
			Stage:           bizdev.MaturityStage(req.Stage),
			ProjectName:     req.ProjectName,
			CustomerName:    req.CustomerName,
			DeviceType:      req.DeviceType,
			IntendedUse:     req.IntendedUse,
			TargetMarket:    req.TargetMarket,
			Fields:          req.Fields,
			Backend:         req.Backend,
			AutoRepair:      req.AutoRepair,
			UserID:          req.UserID,
			ProjectID:       req.ProjectID,
			DocID:           req.DocID,
			AuthorID:        req.AuthorID,
			AuthorName:      req.AuthorName,
			ReviewerID:      req.ReviewerID,
			ReviewerName:    req.ReviewerName,
			ApproverID:      req.ApproverID,
			ApproverName:    req.ApproverName,
			ForceDirectEdit: req.ForceDirectEdit,
		})
		if err != nil {
			streamer.EmitError(err)
		} else {
			streamer.EmitDone("stage_run", time.Since(start))
		}
		return
	}

	if req.TaskKey == "" {
		req.TaskKey = "intake"
	}

	_, err := s.svc.ExecuteTask(ctxWithS, bizdev.TaskRunRequest{
		TaskKey:         req.TaskKey,
		ProjectName:     req.ProjectName,
		CustomerName:    req.CustomerName,
		DeviceType:      req.DeviceType,
		IntendedUse:     req.IntendedUse,
		TargetMarket:    req.TargetMarket,
		Fields:          req.Fields,
		Backend:         req.Backend,
		AutoRepair:      req.AutoRepair,
		UserID:          req.UserID,
		ProjectID:       req.ProjectID,
		DocID:           req.DocID,
		AuthorID:        req.AuthorID,
		AuthorName:      req.AuthorName,
		ReviewerID:      req.ReviewerID,
		ReviewerName:    req.ReviewerName,
		ApproverID:      req.ApproverID,
		ApproverName:    req.ApproverName,
		ForceDirectEdit: req.ForceDirectEdit,
	})
	if err != nil {
		streamer.EmitError(err)
	} else {
		streamer.EmitDone("task_run", time.Since(start))
	}
}