Agent-System/internal/server/task_run_handlers.go

package server

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

	"chat-agent-system/internal/agent"
	"chat-agent-system/internal/bizdev"
	"chat-agent-system/internal/ctxmgr"
	"chat-agent-system/internal/docgov"
	"chat-agent-system/internal/projectinstance"
	"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"`
}

type taskRunResponse struct {
	taskrun.RunRecord
	Resumable bool `json:"resumable"`
}

type resumeTaskRunRequest struct {
	Confirm bool `json:"confirm"`
}

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"
	}
	if req.ProjectID > 0 && req.TaskKey != "" && !s.enforceWorkflowDocTaskAccess(w, r.Context(), req.ProjectID, user.ID, req.TaskKey) {
		return
	}
	req.WorkDir = ""

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

	taskScope, scopeErr := s.resolveTaskMessageScope(r.Context(), req.ProjectID, req.TaskKey)
	if scopeErr != nil {
		writeJSONResponse(w, http.StatusConflict, map[string]string{
			"error_code": "TASK_SCOPE_UNRESOLVED",
			"message":    "无法解析文档任务的工作流阶段作用域",
		})
		return
	}
	if busy, preflightErr := s.preflightTaskDocumentBusy(r.Context(), req); preflightErr != nil {
		http.Error(w, preflightErr.Error(), http.StatusBadRequest)
		return
	} else if busy != nil {
		s.recordBusyPreflightMessages(r.Context(), ctxMgr, req, taskScope, busy)
		writeJSONResponse(w, http.StatusConflict, map[string]any{
			"error":      "document_busy",
			"error_code": "DOC_BUSY",
			"message":    busy.Reason,
			"detail":     busy,
		})
		return
	}

	// 幂等两阶段：先 Claim 占位，仅 owner 写聊天并 Start；loser 不得写消息。
	idemKey := taskrun.BuildTaskRunIdempotencyKey(req.SessionID, req.ClientTurnID)
	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,
		IdempotencyKey:     idemKey,
		TaskKey:            req.TaskKey,
		Stage:              req.Stage,
		WorkflowInstanceID: taskScope.WorkflowInstanceID,
		StageInstanceID:    taskScope.StageInstanceID,
		WorkflowStageKey:   taskScope.StageKey,
		Status:             "pending",
	}
	claimed, created, claimErr := s.taskRuns.ClaimRun(r.Context(), record)
	if claimErr != nil {
		http.Error(w, claimErr.Error(), http.StatusInternalServerError)
		return
	}
	if !created {
		writeJSONResponse(w, http.StatusAccepted, createTaskRunResponse{
			RunID:  claimed.RunID,
			Status: claimed.Status,
		})
		return
	}

	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,
		}
		projectinstance.ApplyScopeToMetadata(userMeta, taskScope)
		_, _ = ctxMgr.RecordUserMessageWithID(r.Context(), req.SessionID, req.UserID, "user", req.UserMessage, "TRIGGER_TASK", "", userMeta)
		asstMeta := map[string]any{
			"run_id":         claimed.RunID,
			"project_id":     req.ProjectID,
			"client_turn_id": req.ClientTurnID,
			"task_key":       req.TaskKey,
			"task_status":    "running",
			"task_request": map[string]any{
				"task_key":       req.TaskKey,
				"project_name":   req.ProjectName,
				"customer_name":  req.CustomerName,
				"device_type":    req.DeviceType,
				"intended_use":   req.IntendedUse,
				"target_market":  req.TargetMarket,
				"session_id":     req.SessionID,
				"user_id":        req.UserID,
				"project_id":     req.ProjectID,
				"client_turn_id": req.ClientTurnID,
				"user_message":   req.UserMessage,
				"doc_id":         req.DocID,
			},
		}
		projectinstance.ApplyScopeToMetadata(asstMeta, taskScope)
		assistantMsgID, _ = ctxMgr.RecordUserMessageWithID(r.Context(), req.SessionID, 0, "assistant", "任务执行中…", "TRIGGER_TASK", "", asstMeta)
		claimed.AssistantMsgID = assistantMsgID
		_ = s.taskRuns.BindAssistantMsgID(r.Context(), claimed.RunID, assistantMsgID)
	}

	taskReq := s.buildBizdevTaskRequest(req.TaskRunRequestDTO, req.SessionID, req.ClientTurnID, claimed.RunID)
	resumeDir := req.WorkDir
	stageReq := s.buildBizdevStageRequest(req.TaskRunRequestDTO, req.SessionID, req.ClientTurnID, claimed.RunID)
	createdRun, err := s.taskRuns.StartClaimedRun(r.Context(), claimed, taskReq, stageReq, resumeDir)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

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

func (s *Server) recordBusyPreflightMessages(
	ctx context.Context,
	ctxMgr *ctxmgr.ContextManager,
	req createTaskRunRequest,
	taskScope *projectinstance.MessageScope,
	busy *docgov.DocumentBusyError,
) {
	if ctxMgr == nil || busy == nil || req.UserMessage == "" {
		return
	}
	userMeta := map[string]any{
		"project_id":     req.ProjectID,
		"client_turn_id": req.ClientTurnID,
		"task_key":       req.TaskKey,
	}
	projectinstance.ApplyScopeToMetadata(userMeta, taskScope)
	_, _ = ctxMgr.RecordUserMessageWithID(ctx, req.SessionID, req.UserID, "user", req.UserMessage, "TRIGGER_TASK", "", userMeta)

	taskRequest := map[string]any{
		"task_key":       req.TaskKey,
		"project_name":   req.ProjectName,
		"customer_name":  req.CustomerName,
		"device_type":    req.DeviceType,
		"intended_use":   req.IntendedUse,
		"target_market":  req.TargetMarket,
		"session_id":     req.SessionID,
		"user_id":        req.UserID,
		"project_id":     req.ProjectID,
		"client_turn_id": req.ClientTurnID,
		"user_message":   req.UserMessage,
		"doc_id":         req.DocID,
	}
	errorInfo := map[string]any{
		"code":         "DOC_BUSY",
		"class":        "fatal",
		"user_message": busy.Reason,
		"actions":      []string{},
		"session_id":   req.SessionID,
	}
	asstMeta := map[string]any{
		"project_id":     req.ProjectID,
		"client_turn_id": req.ClientTurnID,
		"task_key":       req.TaskKey,
		"task_status":    "failed",
		"task_request":   taskRequest,
		"error_info":     errorInfo,
	}
	projectinstance.ApplyScopeToMetadata(asstMeta, taskScope)
	_, _ = ctxMgr.RecordUserMessageWithID(ctx, req.SessionID, 0, "assistant", busy.Reason, "TRIGGER_TASK", "", asstMeta)
}

// preflightTaskDocumentBusy 在写入聊天占位与 task_runs 之前拦截已知租约占用。
// 真正的租约获取仍由 docgov 原子执行；本检查用于避免用户重复点击产生必然失败的 run。
func (s *Server) preflightTaskDocumentBusy(ctx context.Context, req createTaskRunRequest) (*docgov.DocumentBusyError, error) {
	if req.Stage != "" || req.ProjectID <= 0 || req.TaskKey == "" || s.svc == nil || s.svc.DocGov() == nil {
		return nil, nil
	}
	taskReq := s.buildBizdevTaskRequest(req.TaskRunRequestDTO, req.SessionID, req.ClientTurnID, "")
	docID, err := s.svc.ResolveAndValidateDocID(ctx, taskReq)
	if err != nil {
		return nil, err
	}
	if docID == "" {
		return nil, nil
	}
	doc, err := s.svc.DocGov().GetDocument(ctx, docID)
	if errors.Is(err, docgov.ErrDocumentNotFound) {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}
	if !doc.HasActiveLease(time.Now()) {
		return nil, nil
	}
	// 已结束的小队 run 由 AcquireLease 的孤儿回收逻辑处理，不在这里误拦截。
	if doc.LeaseHolderType == docgov.LeaseHolderSquad && s.taskRuns != nil {
		if run, runErr := s.taskRuns.GetRun(doc.LeaseHolderID); runErr == nil && run.Status != "running" {
			return nil, nil
		}
	}
	expiresAt := time.Time{}
	if doc.LeaseExpiresAt != nil {
		expiresAt = *doc.LeaseExpiresAt
	}
	reason := fmt.Sprintf("文档【%s】正在由 %s(%s) 编辑中，当前请求未创建新任务", doc.Title, doc.LeaseHolderType, doc.LeaseHolderID)
	return &docgov.DocumentBusyError{
		DocID:      doc.DocID,
		HolderType: doc.LeaseHolderType,
		HolderID:   doc.LeaseHolderID,
		WorkDir:    doc.LeaseWorkDir,
		ExpiresAt:  expiresAt,
		Reason:     reason,
	}, nil
}

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.assertTaskRunAccess(w, r.Context(), runID, user, taskRunCancel) {
		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")
	agent.ApplyRunnerAuth(req, s.cfg.PiAgentRunnerAuthToken)
	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, s.buildTaskRunResponse(rec))
}

func (s *Server) handleResumeTaskRun(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.assertTaskRunAccess(w, r.Context(), runID, user, taskRunResume) {
		return
	}

	var req resumeTaskRunRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "invalid request body: "+err.Error(), http.StatusBadRequest)
		return
	}
	if !req.Confirm {
		writeJSONResponse(w, http.StatusBadRequest, map[string]string{
			"error_code": "ERR_CONFIRM_REQUIRED",
			"message":    "续跑 interrupted 任务需要 confirm=true",
		})
		return
	}

	rec, err := s.taskRuns.GetRun(runID)
	if err != nil {
		http.Error(w, err.Error(), http.StatusNotFound)
		return
	}
	if rec.Status != "interrupted" {
		writeJSONResponse(w, http.StatusConflict, map[string]string{
			"error_code": "ERR_NOT_RESUMABLE",
			"message":    "仅 interrupted 状态且存在 workspace 的任务可续跑",
		})
		return
	}
	if rec.WorkspaceID == "" {
		writeJSONResponse(w, http.StatusConflict, map[string]string{
			"error_code": "ERR_NO_WORKSPACE",
			"message":    "任务缺少可恢复的工作区",
		})
		return
	}
	if s.resolver == nil {
		http.Error(w, "workspace resolver unavailable", http.StatusServiceUnavailable)
		return
	}
	workDir, werr := s.resolver.ResolveStoredWorkDir(rec.WorkspaceID)
	if werr != nil {
		writeJSONResponse(w, http.StatusForbidden, map[string]string{
			"error_code": "ERR_WORKSPACE_FORBIDDEN",
			"message":    "工作区路径无效或越界",
		})
		return
	}
	_ = workDir // resume path validated; executor receives stored workspace id

	taskReq := s.buildBizdevTaskRequest(TaskRunRequestDTO{
		UserID:    rec.UserID,
		ProjectID: rec.ProjectID,
		TaskKey:   rec.TaskKey,
	}, rec.SessionID, rec.ClientTurnID, rec.RunID)
	stageReq := s.buildBizdevStageRequest(TaskRunRequestDTO{
		UserID:    rec.UserID,
		ProjectID: rec.ProjectID,
		Stage:     rec.Stage,
	}, rec.SessionID, rec.ClientTurnID, rec.RunID)

	resumed, err := s.taskRuns.ResumeInterrupted(r.Context(), rec, rec.WorkspaceID, taskReq, stageReq)
	if err != nil {
		if errors.Is(err, taskrun.ErrRunNotResumable) || errors.Is(err, taskrun.ErrRunAlreadyActive) {
			writeJSONResponse(w, http.StatusConflict, map[string]string{
				"error_code": "ERR_NOT_RESUMABLE",
				"message":    err.Error(),
			})
			return
		}
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

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

func (s *Server) buildTaskRunResponse(rec *taskrun.RunRecord) taskRunResponse {
	resumable := rec != nil && rec.Status == "interrupted" && rec.WorkspaceID != ""
	if resumable && s.resolver != nil {
		if _, err := s.resolver.ResolveStoredWorkDir(rec.WorkspaceID); err != nil {
			resumable = false
		}
	}
	if rec == nil {
		return taskRunResponse{}
	}
	return taskRunResponse{RunRecord: *rec, Resumable: resumable}
}

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)
	})
}