Agent-System/internal/server/task_run_handlers.go

package server

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"strconv"
	"strings"
	"time"

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

type createTaskRunRequest struct {
	TaskRunRequestDTO
	SessionID    string `json:"session_id"`
	ClientTurnID string `json:"client_turn_id"`
	UserMessage  string `json:"user_message"`
	WorkDir      string `json:"work_dir,omitempty"`
}

type createTaskRunResponse struct {
	RunID       string `json:"run_id"`
	WorkspaceID string `json:"workspace_id,omitempty"`
	Status      string `json:"status"`
}

func (s *Server) handleCreateTaskRun(w http.ResponseWriter, r *http.Request) {
	user, ok := s.requirePrincipal(w, r)
	if !ok {
		return
	}
	var req createTaskRunRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "invalid request body: "+err.Error(), http.StatusBadRequest)
		return
	}
	req.UserID = user.ID
	if req.ProjectID > 0 {
		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 err := bizdev.ValidateTaskRunIdentity(req.SessionID, req.ProjectID); err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	if req.TaskKey == "" && req.Stage == "" {
		req.TaskKey = "intake"
	}
	req.WorkDir = ""

	ctxMgr := s.svc.ContextManager()
	if ctxMgr != nil {
		_, _ = ctxMgr.EnsureChatSession(r.Context(), req.SessionID, req.UserID, req.ProjectID, req.ProjectName)
	}

	var assistantMsgID int64
	if ctxMgr != nil && req.UserMessage != "" {
		userMeta := map[string]any{
			"project_id":     req.ProjectID,
			"client_turn_id": req.ClientTurnID,
			"task_key":       req.TaskKey,
		}
		_, _ = ctxMgr.RecordUserMessageWithID(r.Context(), req.SessionID, req.UserID, "user", req.UserMessage, "TRIGGER_TASK", "", userMeta)
		asstMeta := map[string]any{
			"project_id":     req.ProjectID,
			"client_turn_id": req.ClientTurnID,
			"task_key":       req.TaskKey,
			"task_status":    "running",
		}
		assistantMsgID, _ = ctxMgr.RecordUserMessageWithID(r.Context(), req.SessionID, 0, "assistant", "任务执行中…", "TRIGGER_TASK", "", asstMeta)
	}

	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,
		AssistantMsgID: assistantMsgID,
		Status:         "running",
	}

	taskReq := s.buildBizdevTaskRequest(req.TaskRunRequestDTO, req.SessionID, req.ClientTurnID, runID)
	resumeDir := req.WorkDir

	stageReq := s.buildBizdevStageRequest(req.TaskRunRequestDTO, req.SessionID, req.ClientTurnID, runID)
	created, err := s.taskRuns.CreateAndStart(r.Context(), runID, record, taskReq, stageReq, resumeDir)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	writeJSONResponse(w, http.StatusAccepted, createTaskRunResponse{
		RunID:  created.RunID,
		Status: created.Status,
	})
}

func (s *Server) handleTaskRunEvents(w http.ResponseWriter, r *http.Request) {
	user, ok := s.requirePrincipal(w, r)
	if !ok {
		return
	}
	runID := r.PathValue("run_id")
	if runID == "" {
		http.Error(w, "run_id required", http.StatusBadRequest)
		return
	}
	if !s.assertTaskRunOwnedByUser(w, r.Context(), runID, user) {
		return
	}
	flusher, ok := w.(http.Flusher)
	if !ok {
		http.Error(w, "streaming unsupported", http.StatusInternalServerError)
		return
	}

	afterSeq := int64(0)
	if raw := r.URL.Query().Get("after_seq"); raw != "" {
		afterSeq, _ = strconv.ParseInt(raw, 10, 64)
	}

	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")
	flusher.Flush()

	ch, replay, err := s.taskRuns.Subscribe(runID, afterSeq)
	if err != nil {
		http.Error(w, err.Error(), http.StatusNotFound)
		return
	}
	defer s.taskRuns.Unsubscribe(runID, ch)

	writeStoredEvents := func(events []taskrun.StoredEvent) {
		for _, stored := range events {
			data, _ := json.Marshal(stored.Event)
			fmt.Fprintf(w, "id: %d\n", stored.Seq)
			fmt.Fprintf(w, "event: %s\ndata: %s\n\n", stored.Event.Type, string(data))
		}
		flusher.Flush()
	}

	writeStoredEvents(replay)

	for {
		select {
		case <-r.Context().Done():
			return
		case stored, ok := <-ch:
			if !ok {
				return
			}
			writeStoredEvents([]taskrun.StoredEvent{stored})
			if stored.Event.Type == stream.EventDone || stored.Event.Type == stream.EventError {
				return
			}
		}
	}
}

func (s *Server) handleCancelTaskRun(w http.ResponseWriter, r *http.Request) {
	user, ok := s.requirePrincipal(w, r)
	if !ok {
		return
	}
	runID := r.PathValue("run_id")
	if runID == "" {
		http.Error(w, "run_id required", http.StatusBadRequest)
		return
	}
	if !s.assertTaskRunOwnedByUser(w, r.Context(), runID, user) {
		return
	}
	if err := s.taskRuns.Cancel(runID); err != nil {
		http.Error(w, err.Error(), http.StatusNotFound)
		return
	}
	if rec, getErr := s.taskRuns.GetRun(runID); getErr == nil && rec.WorkspaceID != "" {
		s.cancelPiAgentRunner(rec.WorkspaceID)
	}
	writeJSONResponse(w, http.StatusOK, map[string]string{"run_id": runID, "status": "cancelled"})
}

func (s *Server) cancelPiAgentRunner(workspaceID string) {
	runnerURL := strings.TrimRight(strings.TrimSpace(s.cfg.PiAgentRunnerURL), "/")
	if runnerURL == "" || workspaceID == "" {
		return
	}
	body, _ := json.Marshal(map[string]string{"workspace": workspaceID})
	req, err := http.NewRequest(http.MethodPost, runnerURL+"/v1/cancel", bytes.NewReader(body))
	if err != nil {
		return
	}
	req.Header.Set("Content-Type", "application/json")
	client := &http.Client{Timeout: 5 * time.Second}
	resp, err := client.Do(req)
	if err != nil {
		return
	}
	_ = resp.Body.Close()
}

func (s *Server) handleGetTaskRun(w http.ResponseWriter, r *http.Request) {
	user, ok := s.requirePrincipal(w, r)
	if !ok {
		return
	}
	runID := r.PathValue("run_id")
	if !s.assertTaskRunOwnedByUser(w, r.Context(), runID, user) {
		return
	}
	rec, err := s.taskRuns.GetRun(runID)
	if err != nil {
		http.Error(w, err.Error(), http.StatusNotFound)
		return
	}
	writeJSONResponse(w, http.StatusOK, rec)
}

func (s *Server) buildBizdevTaskRequest(dto TaskRunRequestDTO, sessionID, clientTurnID, runID string) bizdev.TaskRunRequest {
	if dto.TaskKey == "" {
		dto.TaskKey = "intake"
	}
	return bizdev.TaskRunRequest{
		TaskKey:         dto.TaskKey,
		ProjectName:     dto.ProjectName,
		CustomerName:    dto.CustomerName,
		DeviceType:      dto.DeviceType,
		IntendedUse:     dto.IntendedUse,
		TargetMarket:    dto.TargetMarket,
		Fields:          dto.Fields,
		Materials:       dto.Materials,
		Backend:         dto.Backend,
		AutoRepair:      dto.AutoRepair,
		SessionID:       sessionID,
		UserID:          dto.UserID,
		ProjectID:       dto.ProjectID,
		ClientTurnID:    clientTurnID,
		RunID:           runID,
		DocID:           dto.DocID,
		AuthorID:        dto.AuthorID,
		AuthorName:      dto.AuthorName,
		ReviewerID:      dto.ReviewerID,
		ReviewerName:    dto.ReviewerName,
		ApproverID:      dto.ApproverID,
		ApproverName:    dto.ApproverName,
		ForceDirectEdit: dto.ForceDirectEdit,
		CustomSteps:     dto.CustomSteps,
	}
}

func (s *Server) buildBizdevStageRequest(dto TaskRunRequestDTO, sessionID, clientTurnID, runID string) bizdev.StageRunRequest {
	return bizdev.StageRunRequest{
		Stage:           bizdev.MaturityStage(dto.Stage),
		ProjectName:     dto.ProjectName,
		CustomerName:    dto.CustomerName,
		DeviceType:      dto.DeviceType,
		IntendedUse:     dto.IntendedUse,
		TargetMarket:    dto.TargetMarket,
		Fields:          dto.Fields,
		Materials:       dto.Materials,
		Backend:         dto.Backend,
		AutoRepair:      dto.AutoRepair,
		SessionID:       sessionID,
		UserID:          dto.UserID,
		ProjectID:       dto.ProjectID,
		ClientTurnID:    clientTurnID,
		RunID:           runID,
		DocID:           dto.DocID,
		AuthorID:        dto.AuthorID,
		AuthorName:      dto.AuthorName,
		ReviewerID:      dto.ReviewerID,
		ReviewerName:    dto.ReviewerName,
		ApproverID:      dto.ApproverID,
		ApproverName:    dto.ApproverName,
		ForceDirectEdit: dto.ForceDirectEdit,
	}
}

func (s *Server) initTaskRunExecutors() {
	s.taskRuns.SetExecutor(
		func(ctx context.Context, _ string, req bizdev.TaskRunRequest, _ *stream.Streamer) (*bizdev.TaskRunResult, error) {
			return s.svc.ExecuteTask(ctx, req)
		},
		func(ctx context.Context, _ string, workDir string, _ *stream.Streamer) (*bizdev.TaskRunResult, error) {
			res, err := s.svc.ResumeOrModifyWorkspace(ctx, squad.ModifyWorkflowRequest{WorkDir: workDir})
			if res == nil {
				return nil, err
			}
			out := squadResultToTaskRun(res, res.TaskKey)
			if out != nil && out.WorkDir == "" {
				out.WorkDir = workDir
			}
			return out, err
		},
		func(ctx context.Context, _ string, req bizdev.StageRunRequest, _ *stream.Streamer) (*bizdev.TaskRunResult, error) {
			stageRes, err := s.svc.ExecuteStage(ctx, req)
			if stageRes == nil {
				return nil, err
			}
			if stageRes.PrimaryArtifact != nil {
				return stageRes.PrimaryArtifact, err
			}
			return &bizdev.TaskRunResult{
				WorkDir:       stageRes.WorkDir,
				Success:       stageRes.Success,
				ErrorReason:   stageRes.ErrorReason,
				ArtifactTitle: stageRes.StageTitle,
			}, err
		},
	)
	s.taskRuns.SetFinalizer(func(ctx context.Context, record *taskrun.RunRecord, result *bizdev.TaskRunResult) (string, error) {
		return s.finalizeTaskRun(ctx, record, result)
	})
}