Files
Agent-System/internal/server/sse.go
T

308 lines
9.5 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
package server
import (
"encoding/json"
"fmt"
"net/http"
"strconv"
"sync"
"time"
"chat-agent-system/internal/bizdev"
"chat-agent-system/internal/squad"
"chat-agent-system/internal/stream"
)
// 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"`
}
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 响应头 + 请求追踪 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")
w.Header().Set("Access-Control-Allow-Origin", "*")
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,
)
if req.UserID > 0 {
sess.UserID = req.UserID
}
if req.ProjectID > 0 {
sess.ProjectID = req.ProjectID
}
// 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(ctx, req.Message, listener, req.IncludeDetails)
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")
w.Header().Set("Access-Control-Allow-Origin", "*")
w.Header().Set("X-Accel-Buffering", "no") // 禁用 Nginx 等反代缓冲
flusher.Flush()
var writeMu sync.Mutex
// lastPhaseID 记录最近一次 phase 事件,失败时用于 error_info.phase_id。
// Streamer.Emit 为同步分发,listener 与后续读取同处 handler goroutine。
lastPhaseID := ""
listener := func(ev stream.StreamEvent) {
writeMu.Lock()
defer writeMu.Unlock()
if ev.Type == stream.EventPhase && ev.Phase != nil && ev.Phase.PhaseID != "" {
lastPhaseID = ev.Phase.PhaseID
}
data, err := json.Marshal(ev)
if err != nil {
return
}
fmt.Fprintf(w, "event: %s\ndata: %s\n\n", ev.Type, string(data))
flusher.Flush()
}
// 分类失败并发射结构化错误:phaseID 为空时回落到流中最近一次 phase
ctxWithS, streamer := stream.WithListener(stream.WithRequestID(r.Context(), requestID), includeDetails, listener)
classifyAndEmitError := func(err error, phaseID, workspaceID string) {
if phaseID == "" {
writeMu.Lock()
phaseID = lastPhaseID
writeMu.Unlock()
}
streamer.EmitErrorInfo(stream.ClassifyError(err, phaseID, requestID, "", workspaceID))
}
start := time.Now()
// resume 模式:work_dir 非空时续跑既有工作区(P2:「继续任务」后端入口),优先于 stage/task 新任务
if req.WorkDir != "" {
res, err := s.svc.ResumeOrModifyWorkspace(ctxWithS, squad.ModifyWorkflowRequest{
WorkDir: req.WorkDir,
AddSteps: req.AddSteps,
Backend: req.Backend,
AutoRepair: req.AutoRepair,
})
if err != nil {
wsID := req.WorkDir
if res != nil && res.WorkDir != "" {
wsID = res.WorkDir
}
classifyAndEmitError(err, "task_resume", wsID)
} else {
streamer.EmitDoneWithWorkspace("task_resume", req.WorkDir, time.Since(start))
}
return
}
if req.Stage != "" {
res, 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 {
wsID := ""
if res != nil && res.WorkDir != "" {
wsID = res.WorkDir
}
classifyAndEmitError(err, "stage_run", wsID)
} else {
wsID := ""
if res != nil {
wsID = res.WorkDir
}
streamer.EmitDoneWithWorkspace("stage_run", wsID, time.Since(start))
}
return
}
if req.TaskKey == "" {
req.TaskKey = "intake"
}
res, 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 {
// 软失败(执行器返回带 WorkDir 的 result,如 docgov 锁定 / squad 层任务中断)→ resumable;
// 硬失败(工作区未建立)→ workspaceID 为空,按超时/网络/默认规则分类
wsID := ""
if res != nil && res.WorkDir != "" {
wsID = res.WorkDir
}
classifyAndEmitError(err, "", wsID)
} else {
wsID := ""
if res != nil {
wsID = res.WorkDir
}
streamer.EmitDoneWithWorkspace("task_run", wsID, time.Since(start))
}
}