Agent-System/internal/server/sse.go

package server

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

	"chat-agent-system/internal/chat"
	"chat-agent-system/internal/ctxmgr"
	"chat-agent-system/internal/squad"
	"chat-agent-system/internal/stream"
	"chat-agent-system/internal/taskrun"
)

// taskStreamRequest 任务/阶段流式运行请求：在 handlers.go 的 TaskRunRequestDTO 基础上
// 增量扩展 resume 能力（work_dir 非空 → 续跑既有工作区），不改动既有 DTO 定义（向后兼容）。
type taskStreamRequest struct {
	TaskRunRequestDTO
	WorkDir  string           `json:"work_dir,omitempty"`
	AddSteps []squad.StepSpec `json:"add_steps,omitempty"`
}

// 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"`
	ClientTurnID   string `json:"client_turn_id,omitempty"`
}

func writeSSE(w http.ResponseWriter, flusher http.Flusher, ev stream.StreamEvent) {
	data, err := json.Marshal(ev)
	if err != nil {
		return
	}
	fmt.Fprintf(w, "event: %s\ndata: %s\n\n", ev.Type, string(data))
	flusher.Flush()
}

func replayExistingTurn(w http.ResponseWriter, flusher http.Flusher, rec *ctxmgr.ChatTurnRecord) {
	status := string(rec.Status)
	duplicate := true
	if rec.UserMessageID > 0 {
		writeSSE(w, flusher, stream.StreamEvent{
			Type:      stream.EventTurnAck,
			Timestamp: time.Now(),
			Turn: &stream.TurnAckPayload{
				ClientTurnID:  rec.ClientTurnID,
				UserMessageID: rec.UserMessageID,
				Status:        status,
				Duplicate:     duplicate,
			},
		})
	}
	switch rec.Status {
	case ctxmgr.TurnStatusCompleted:
		if rec.AssistantMessageID > 0 {
			writeSSE(w, flusher, stream.StreamEvent{
				Type:      stream.EventDone,
				Timestamp: time.Now(),
				Done: &stream.DonePayload{
					SessionID:     rec.SessionID,
					ClientTurnID:  rec.ClientTurnID,
					UserMessageID: rec.UserMessageID,
					MessageID:     rec.AssistantMessageID,
					Status:        "completed",
				},
			})
		}
	case ctxmgr.TurnStatusFailed, ctxmgr.TurnStatusCancelled:
		writeSSE(w, flusher, stream.StreamEvent{
			Type:      stream.EventError,
			Timestamp: time.Now(),
			Error:     "该回合已失败或已取消，请重新发送消息",
			ErrorInfo: &stream.ErrorInfo{
				Code:        "TURN_NOT_REPLAYABLE",
				UserMessage: "该回合已失败或已取消，请重新发送消息",
				SessionID:   rec.SessionID,
			},
		})
	case ctxmgr.TurnStatusProcessing:
		stale := rec.LeaseExpiresAt != nil && time.Now().After(*rec.LeaseExpiresAt)
		if stale {
			writeSSE(w, flusher, stream.StreamEvent{
				Type:      stream.EventError,
				Timestamp: time.Now(),
				Error:     "上一回合处理超时，请重新发送消息",
				ErrorInfo: &stream.ErrorInfo{
					Code:        "TURN_LEASE_EXPIRED",
					UserMessage: "上一回合处理超时，请重新发送消息",
					SessionID:   rec.SessionID,
				},
			})
		}
	}
}

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
	}

	user, ok := s.requirePrincipal(w, r)
	if !ok {
		return
	}
	req.UserID = user.ID

	// 1. 设置 SSE 响应头 + 请求追踪 ID
	requestID := fmt.Sprintf("req_%d", time.Now().UnixNano())
	w.Header().Set("X-Request-Id", requestID)
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	s.writeCORS(w, r)
	w.Header().Set("X-Accel-Buffering", "no") // 禁用 Nginx 等反代缓冲
	flusher.Flush()

	ctx := stream.WithRequestID(r.Context(), requestID)

	// 2. 准备会话
	sess := s.getOrCreateSession(
		req.SessionID,
		req.ProjectName,
		req.CustomerName,
		req.DeviceType,
		req.IntendedUse,
		req.TargetMarket,
		req.WorkDir,
		req.UserName,
		req.UserRole,
	)
	sess.UserID = req.UserID
	if req.ProjectID > 0 {
		sess.ProjectID = req.ProjectID
	}
	if sess.UserCtx != nil {
		if req.ProjectName != "" && (sess.UserCtx.Project.Name == "" || sess.UserCtx.Project.Name == "医疗器械研发项目" || sess.UserCtx.Project.Name == "当前项目") {
			sess.UserCtx.Project.Name = req.ProjectName
		}
		if req.DeviceType != "" {
			sess.UserCtx.Project.DeviceType = req.DeviceType
		}
		if req.UserName != "" {
			sess.UserCtx.User.Name = req.UserName
		}
		if req.UserRole != "" {
			sess.UserCtx.User.Role = 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()
	}

	clientTurnID := req.ClientTurnID
	if clientTurnID == "" {
		clientTurnID = r.Header.Get("Idempotency-Key")
	}

	var turnOpts *chat.TurnOptions
	if clientTurnID != "" {
		turnOpts = &chat.TurnOptions{ClientTurnID: clientTurnID}
		if sess.CtxMgr != nil {
			bodyHash := ctxmgr.HashTurnBody(req.SessionID, req.Message)
			acq, acqErr := sess.CtxMgr.AcquireChatTurn(ctx, req.SessionID, clientTurnID, bodyHash, requestID, sess.UserID, sess.ProjectID)
			if errors.Is(acqErr, ctxmgr.ErrIdempotencyKeyReused) {
				writeMu.Lock()
				info := stream.ClassifyError(acqErr, "idempotency", requestID, sess.ID, "")
				info.Code = "IDEMPOTENCY_KEY_REUSED"
				writeSSE(w, flusher, stream.StreamEvent{
					Type:      stream.EventError,
					Timestamp: time.Now(),
					Error:     info.UserMessage,
					ErrorInfo: &info,
				})
				writeMu.Unlock()
				return
			}
			if acqErr != nil {
				writeMu.Lock()
				info := stream.ClassifyError(acqErr, "turn_acquire", requestID, sess.ID, "")
				writeSSE(w, flusher, stream.StreamEvent{
					Type:      stream.EventError,
					Timestamp: time.Now(),
					Error:     info.UserMessage,
					ErrorInfo: &info,
				})
				writeMu.Unlock()
				return
			}
			if acq != nil && !acq.WonLease {
				writeMu.Lock()
				replayExistingTurn(w, flusher, acq.Record)
				writeMu.Unlock()
				return
			}
		}
	}

	// 4. 同 session 串行执行后发起流式对话
	var err error
	execErr := sess.WithExecLock(func() error {
		var sendErr error
		_, sendErr = sess.SendMessageStreamTurn(ctx, req.Message, listener, req.IncludeDetails, turnOpts)
		return sendErr
	})
	err = execErr
	if err != nil {
		writeMu.Lock()
		info := stream.ClassifyError(err, "chat_stream", requestID, sess.ID, "")
		errEv := stream.StreamEvent{
			Type:      stream.EventError,
			Timestamp: time.Now(),
			Error:     info.UserMessage,
			ErrorInfo: &info,
		}
		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 taskStreamRequest
	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.WorkDir = q.Get("work_dir")
		req.AutoRepair, _ = strconv.ParseBool(q.Get("auto_repair"))
		if inc := q.Get("include_details"); inc != "" {
			includeDetails, _ = strconv.ParseBool(inc)
		}
	}

	// 请求追踪 ID：与 chat stream 对齐，注入 ctx 并回写响应头
	requestID := fmt.Sprintf("req_%d", time.Now().UnixNano())
	w.Header().Set("X-Request-Id", requestID)
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	s.writeCORS(w, r)
	w.Header().Set("X-Accel-Buffering", "no") // 禁用 Nginx 等反代缓冲
	flusher.Flush()

	user, ok := s.requirePrincipal(w, r)
	if !ok {
		return
	}

	if req.TaskKey == "" && req.Stage == "" {
		req.TaskKey = "intake"
	}
	req.UserID = user.ID
	req.WorkDir = ""

	runID := fmt.Sprintf("run_%d", time.Now().UnixNano())
	record := &taskrun.RunRecord{
		RunID:        runID,
		SessionID:    req.SessionID,
		ProjectID:    req.ProjectID,
		UserID:       req.UserID,
		ClientTurnID: req.ClientTurnID,
		TaskKey:      req.TaskKey,
		Stage:        req.Stage,
		Status:       "running",
	}
	taskReq := s.buildBizdevTaskRequest(req.TaskRunRequestDTO, req.SessionID, req.ClientTurnID, runID)
	stageReq := s.buildBizdevStageRequest(req.TaskRunRequestDTO, req.SessionID, req.ClientTurnID, runID)

	created, err := s.taskRuns.CreateAndStart(r.Context(), runID, record, taskReq, stageReq, req.WorkDir)
	if err != nil {
		http.Error(w, "start task run failed: "+err.Error(), http.StatusInternalServerError)
		return
	}
	_ = includeDetails
	_ = requestID

	ch, replay, err := s.taskRuns.Subscribe(created.RunID, 0)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	defer s.taskRuns.Unsubscribe(created.RunID, ch)

	writeStored := func(stored taskrun.StoredEvent) {
		data, marshalErr := json.Marshal(stored.Event)
		if marshalErr != nil {
			return
		}
		fmt.Fprintf(w, "id: %d\n", stored.Seq)
		fmt.Fprintf(w, "event: %s\ndata: %s\n\n", stored.Event.Type, string(data))
		flusher.Flush()
	}

	for _, stored := range replay {
		writeStored(stored)
		if stored.Event.Type == stream.EventDone || stored.Event.Type == stream.EventError {
			return
		}
	}

	for {
		select {
		case <-r.Context().Done():
			// SSE 断线仅停止推送，不取消后台 worker
			return
		case stored, ok := <-ch:
			if !ok {
				return
			}
			writeStored(stored)
			if stored.Event.Type == stream.EventDone || stored.Event.Type == stream.EventError {
				return
			}
		}
	}
}