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"` } // 长任务 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)) } func writeSSEHeartbeat(w http.ResponseWriter, flusher http.Flusher) error { extendSSEWriteDeadline(w) if _, err := fmt.Fprintf(w, ": heartbeat %d\n\n", time.Now().Unix()); err != nil { return err } flusher.Flush() return nil } 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 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 == "" { http.Error(w, "message cannot be empty", http.StatusBadRequest) return } user, ok := s.requirePrincipal(w, r) if !ok { return } req.UserID = user.ID req.WorkDir = "" // 会话与项目绑定必须在 SSE headers/Flush 之前完成授权(与 tasks/run/stream 对齐) 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.ProjectID > 0 { if !s.enforceChatProjectAccess(w, r, user, sess.ProjectID) { 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() ctx := stream.WithRequestID(r.Context(), requestID) s.enrichSessionWorkflowStage(ctx, sess, req.UserID) if sess.ProjectID > 0 && !s.allowActiveStageChatAccess(ctx, sess.ProjectID, 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 监听分发闭包 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: \ndata: \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 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 } } } // 请求追踪 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() 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 } } } }