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/stream/transport"
	"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"`
}

// 长任务 SSE：不取消全局 WriteTimeout，仅在流式响应上周期性延长写期限并发心跳。
const (
	sseWriteDeadlineExtend = 45 * time.Second
	sseHeartbeatInterval   = 15 * time.Second
)

func extendSSEWriteDeadline(w http.ResponseWriter) {
	rc := http.NewResponseController(w)
	_ = rc.SetWriteDeadline(time.Now().Add(sseWriteDeadlineExtend))
}

// syncTaskWriteDeadline 覆盖同步 /tasks/run 的长执行窗口。全局 WriteTimeout=120s
// 会在响应写出前截断数十分钟的阶段/任务；此处只延长该请求的写期限，不改全局值。
const syncTaskWriteDeadline = 75 * time.Minute

func extendSyncTaskWriteDeadline(w http.ResponseWriter) {
	rc := http.NewResponseController(w)
	_ = rc.SetWriteDeadline(time.Now().Add(syncTaskWriteDeadline))
}

func writeSSEHeartbeat(w http.ResponseWriter, flusher http.Flusher) error {
	extendSSEWriteDeadline(w)
	if err := (transport.Encoder{W: w}).Heartbeat(time.Now()); err != nil {
		return err
	}
	flusher.Flush()
	return nil
}

func writeSSE(w http.ResponseWriter, flusher http.Flusher, ev stream.StreamEvent) {
	_ = (transport.Encoder{W: w}).Write(ev)
	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 {
		writeSafeError(w, r, http.StatusInternalServerError, "streaming unsupported by underlying connection")
		return
	}

	var req StreamChatRequest
	if r.Method == http.MethodPost {
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			writeSafeError(w, r, http.StatusBadRequest, "invalid request body: "+err.Error())
			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 pid := q.Get("project_id"); pid != "" {
			req.ProjectID, _ = strconv.ParseInt(pid, 10, 64)
		}
		if inc := q.Get("include_details"); inc != "" {
			req.IncludeDetails, _ = strconv.ParseBool(inc)
		}
	}

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

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

	// 会话与项目绑定必须在 SSE headers/Flush 之前完成授权（与 tasks/run/stream 对齐）
	sess, err := s.getOrCreateSession(
		r.Context(),
		user.ID,
		req.SessionID,
		req.ProjectName,
		req.CustomerName,
		req.DeviceType,
		req.IntendedUse,
		req.TargetMarket,
		req.WorkDir,
		req.UserName,
		req.UserRole,
	)
	if err != nil {
		if errors.Is(err, ctxmgr.ErrSessionOwnerMismatch) {
			writeJSONResponse(w, http.StatusForbidden, map[string]string{
				"error_code": "SESSION_FORBIDDEN",
				"message":    "无权访问该会话",
			})
			return
		}
		writeSafeError(w, r, http.StatusInternalServerError, "resolve session failed: "+err.Error())
		return
	}
	s.hydrateVisibleHistory(r.Context(), sess, user.ID)
	if req.ProjectID > 0 && req.ProjectID != sess.Project() {
		sess.SetProjectID(req.ProjectID)
	}
	if !s.enforceChatRequiresProject(w, r, user, sess.Project()) {
		return
	}
	if sess.Project() > 0 {
		if !s.enforceChatProjectAccess(w, r, user, sess.Project()) {
			return
		}
	}

	// 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()
	extendSSEWriteDeadline(w)

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

	s.enrichSessionWorkflowStage(ctx, sess, req.UserID)
	if sess.Project() > 0 && !s.allowActiveStageChatAccess(ctx, sess.Project(), req.UserID) {
		writeWorkflowAccessDeniedSSE(w, flusher, req.SessionID)
		return
	}
	s.bindSessionDocumentSummaryReader(sess, user)
	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 监听分发闭包；心跳与业务写共用 writeMu，避免与 WriteTimeout 竞态
	var writeMu sync.Mutex
	listener := func(ev stream.StreamEvent) {
		writeMu.Lock()
		defer writeMu.Unlock()
		extendSSEWriteDeadline(w)
		writeSSE(w, flusher, ev)
	}

	heartbeatDone := make(chan struct{})
	var heartbeatWg sync.WaitGroup
	heartbeatWg.Add(1)
	go func() {
		defer heartbeatWg.Done()
		ticker := time.NewTicker(sseHeartbeatInterval)
		defer ticker.Stop()
		for {
			select {
			case <-heartbeatDone:
				return
			case <-r.Context().Done():
				return
			case <-ticker.C:
				writeMu.Lock()
				err := writeSSEHeartbeat(w, flusher)
				writeMu.Unlock()
				if err != nil {
					return
				}
			}
		}
	}()
	// handler 返回前先停心跳再等待其退出，避免心跳 goroutine 在连接回收后残留写
	defer func() {
		close(heartbeatDone)
		heartbeatWg.Wait()
	}()

	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 串行执行后发起流式对话
	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, "")
		writeSSE(w, flusher, stream.StreamEvent{
			Type:      stream.EventError,
			Timestamp: time.Now(),
			Error:     info.UserMessage,
			ErrorInfo: &info,
		})
		writeMu.Unlock()
	}
}

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

	var req taskStreamRequest
	includeDetails := false
	if r.Method == http.MethodPost {
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			writeSafeError(w, r, http.StatusBadRequest, "invalid request body: "+err.Error())
			return
		}
	} 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 = resolveTaskBackend(s, q.Get("backend"))
		req.WorkDir = q.Get("work_dir")
		req.AutoRepair, _ = strconv.ParseBool(q.Get("auto_repair"))
		if pid := q.Get("project_id"); pid != "" {
			req.ProjectID, _ = strconv.ParseInt(pid, 10, 64)
		}
		if inc := q.Get("include_details"); inc != "" {
			includeDetails, _ = strconv.ParseBool(inc)
		}
	}

	user, ok := s.requirePrincipal(w, r)
	if !ok {
		return
	}
	if req.ProjectID > 0 {
		if !s.isSystemAdminUser(r.Context(), user) {
			if err := s.authorizer.RequireProjectAccess(r.Context(), req.ProjectID, user.ID); err != nil {
				writeJSONResponse(w, http.StatusForbidden, map[string]string{
					"error_code": "ERR_FORBIDDEN",
					"message":    "无权访问该项目",
				})
				return
			}
		}
	}

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

	// 请求追踪 ID：与 chat stream 对齐；CreateAndStart/Subscribe 必须在 Flush 之前，
	// 失败时返回普通 JSON 错误，避免污染已声明的 text/event-stream 连接。
	requestID := fmt.Sprintf("req_%d", time.Now().UnixNano())
	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)

	// 幂等前置：与 REST 路径对齐，EventSource 断线重连命中已有 run 时直接复用，
	// 避免重复执行同一 client_turn_id 的任务。
	var reused *taskrun.RunRecord
	if idemKey := taskrun.BuildTaskRunIdempotencyKey(req.SessionID, req.ClientTurnID); idemKey != "" && s.taskRuns != nil {
		if existing, findErr := s.taskRuns.FindByIdempotencyKey(r.Context(), idemKey); findErr != nil {
			writeSafeError(w, r, http.StatusInternalServerError, findErr.Error())
			return
		} else if existing != nil {
			if existing.UserID != req.UserID || existing.SessionID != req.SessionID ||
				existing.ProjectID != req.ProjectID || existing.ClientTurnID != req.ClientTurnID {
				writeJSONResponse(w, http.StatusForbidden, map[string]string{
					"error_code": "ERR_FORBIDDEN",
					"message":    "幂等键已存在但不属于当前用户/会话/项目",
				})
				return
			}
			reused = existing
		}
	}

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

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

	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()
	extendSSEWriteDeadline(w)

	writeStored := func(stored taskrun.StoredEvent) {
		_ = (transport.Encoder{W: w}).WriteWithID(stored.Seq, stored.Event)
		flusher.Flush()
		extendSSEWriteDeadline(w)
	}

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

	heartbeat := time.NewTicker(sseHeartbeatInterval)
	defer heartbeat.Stop()

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