Agent-System/internal/taskrun/manager.go

package taskrun

import (
	"context"
	"fmt"
	"log/slog"
	"sync"
	"time"

	"chat-agent-system/internal/bizdev"
	"chat-agent-system/internal/ctxmgr"
	"chat-agent-system/internal/projectworkflow"
	"chat-agent-system/internal/stream"
)

// Executor 执行具体任务逻辑（由 server 注入）。
type Executor func(ctx context.Context, runID string, req bizdev.TaskRunRequest, streamer *stream.Streamer) (*bizdev.TaskRunResult, error)

// ResumeExecutor 续跑既有工作区。
type ResumeExecutor func(ctx context.Context, runID string, workDir string, streamer *stream.Streamer) (*bizdev.TaskRunResult, error)

// StageExecutor 执行成熟度阶段任务序列。
type StageExecutor func(ctx context.Context, runID string, req bizdev.StageRunRequest, streamer *stream.Streamer) (*bizdev.TaskRunResult, error)

// Finalizer 任务成功后投影运行记录与聊天消息（由 server 注入）。
type Finalizer func(ctx context.Context, record *RunRecord, result *bizdev.TaskRunResult) (terminalStatus string, err error)

// Manager 后台任务运行观察者：SSE 断线不取消 worker。
type Manager struct {
	mu       sync.RWMutex
	runs     map[string]*activeRun
	exec     Executor
	resume   ResumeExecutor
	stage    StageExecutor
	finalize Finalizer
	pwSvc    *projectworkflow.Service
	ctxMgr   *ctxmgr.ContextManager
}

const runRetention = time.Hour

func NewManager(pwSvc *projectworkflow.Service, ctxMgr *ctxmgr.ContextManager) *Manager {
	m := &Manager{
		runs:   make(map[string]*activeRun),
		pwSvc:  pwSvc,
		ctxMgr: ctxMgr,
	}
	go m.cleanupLoop()
	return m
}

func (m *Manager) cleanupLoop() {
	ticker := time.NewTicker(10 * time.Minute)
	defer ticker.Stop()
	for range ticker.C {
		m.evictStaleRuns()
	}
}

func (m *Manager) evictStaleRuns() {
	cutoff := time.Now().Add(-runRetention)
	m.mu.Lock()
	defer m.mu.Unlock()
	for id, ar := range m.runs {
		ar.mu.RLock()
		status := ar.record.Status
		updated := ar.record.UpdatedAt
		ar.mu.RUnlock()
		if updated.Before(cutoff) && status != "running" {
			delete(m.runs, id)
		}
	}
}

func (m *Manager) SetExecutor(exec Executor, resume ResumeExecutor, stage StageExecutor) {
	m.exec = exec
	m.resume = resume
	m.stage = stage
}

func (m *Manager) SetFinalizer(f Finalizer) {
	m.finalize = f
}

func (m *Manager) CreateAndStart(parent context.Context, runID string, record *RunRecord, taskReq bizdev.TaskRunRequest, stageReq bizdev.StageRunRequest, resumeWorkDir string) (*RunRecord, error) {
	if runID == "" {
		runID = fmt.Sprintf("run_%d", time.Now().UnixNano())
	}
	if record == nil {
		record = &RunRecord{RunID: runID}
	}
	record.RunID = runID
	if record.Status == "" {
		record.Status = "running"
	}
	now := time.Now()
	record.CreatedAt = now
	record.UpdatedAt = now

	workerCtx, cancel := context.WithCancel(context.WithoutCancel(parent))

	ar := &activeRun{
		record: record,
		subs:   make(map[int]chan StoredEvent),
		cancel: cancel,
		done:   make(chan struct{}),
	}

	m.mu.Lock()
	m.runs[runID] = ar
	m.mu.Unlock()

	if m.pwSvc != nil && record.ProjectID > 0 {
		startedBy := record.UserID
		_, _ = m.pwSvc.CreateWorkflowRun(workerCtx, &projectworkflow.WorkflowRun{
			RunID:     runID,
			ProjectID: record.ProjectID,
			Status:    "running",
			StartedBy: &startedBy,
			StageKey:  record.Stage,
			TaskKey:   record.TaskKey,
			InputManifest: map[string]any{
				"session_id":       record.SessionID,
				"client_turn_id":   record.ClientTurnID,
				"assistant_msg_id": record.AssistantMsgID,
			},
		})
	}

	go m.runWorker(workerCtx, ar, taskReq, stageReq, resumeWorkDir)
	return record, nil
}

func (m *Manager) runWorker(ctx context.Context, ar *activeRun, taskReq bizdev.TaskRunRequest, stageReq bizdev.StageRunRequest, resumeWorkDir string) {
	defer close(ar.done)
	taskReq.RunID = ar.record.RunID
	taskReq.SessionID = ar.record.SessionID
	taskReq.ProjectID = ar.record.ProjectID
	taskReq.UserID = ar.record.UserID
	taskReq.ClientTurnID = ar.record.ClientTurnID

	lastPhaseID := ""
	listener := func(ev stream.StreamEvent) {
		if ev.Type == stream.EventPhase && ev.Phase != nil && ev.Phase.PhaseID != "" {
			lastPhaseID = ev.Phase.PhaseID
		}
		m.appendEvent(ar, ev)
	}
	ctxWithS, streamer := stream.WithListener(ctx, false, listener)

	var taskResult *bizdev.TaskRunResult
	var execErr error
	start := time.Now()

	if ar.record.Stage != "" && m.stage != nil {
		stageReq.RunID = ar.record.RunID
		stageReq.SessionID = ar.record.SessionID
		stageReq.ProjectID = ar.record.ProjectID
		stageReq.UserID = ar.record.UserID
		stageReq.ClientTurnID = ar.record.ClientTurnID
		taskResult, execErr = m.stage(ctxWithS, ar.record.RunID, stageReq, streamer)
	} else if resumeWorkDir != "" && m.resume != nil {
		taskResult, execErr = m.resume(ctxWithS, ar.record.RunID, resumeWorkDir, streamer)
	} else if m.exec != nil {
		taskResult, execErr = m.exec(ctxWithS, ar.record.RunID, taskReq, streamer)
	} else {
		execErr = fmt.Errorf("task executor not configured")
	}

	workspaceID := ""
	if taskResult != nil && taskResult.WorkDir != "" {
		workspaceID = taskResult.WorkDir
	} else if resumeWorkDir != "" {
		workspaceID = resumeWorkDir
	}

	if workspaceID != "" {
		ar.mu.Lock()
		ar.record.WorkspaceID = workspaceID
		ar.mu.Unlock()
	}

	if execErr != nil {
		if ctx.Err() == context.Canceled {
			m.finishRun(ar, "cancelled", workspaceID, execErr.Error(), nil)
			info := stream.ClassifyError(fmt.Errorf("task cancelled"), "task_cancel", ar.record.RunID, ar.record.SessionID, workspaceID)
			m.appendEvent(ar, stream.StreamEvent{Type: stream.EventError, ErrorInfo: &info})
			return
		}
		classifyWS := workspaceID
		if classifyWS == "" && resumeWorkDir != "" {
			classifyWS = resumeWorkDir
		}
		info := stream.ClassifyError(execErr, lastPhaseID, ar.record.RunID, ar.record.SessionID, classifyWS)
		m.appendEvent(ar, stream.StreamEvent{Type: stream.EventError, ErrorInfo: &info})
		m.finishRun(ar, "failed", workspaceID, execErr.Error(), taskResult)
		return
	}

	m.appendEvent(ar, stream.StreamEvent{
		Type: stream.EventDone,
		Done: &stream.DonePayload{
			SessionID:   "task_run",
			WorkspaceID: workspaceID,
			DurationMs:  time.Since(start).Milliseconds(),
		},
	})

	terminalStatus := "completed"
	if m.finalize != nil && taskResult != nil {
		if ts, finErr := m.finalize(ctx, ar.record, taskResult); finErr != nil {
			m.finishRun(ar, "finalizing_failed", workspaceID, finErr.Error(), taskResult)
			return
		} else if ts != "" {
			terminalStatus = ts
		}
	} else if taskResult != nil && taskResult.QualityReport != nil && !taskResult.QualityReport.Passed {
		terminalStatus = "quality_failed"
	}
	m.finishRun(ar, terminalStatus, workspaceID, "", taskResult)
}

func (m *Manager) finishRun(ar *activeRun, status, workspaceID, errMsg string, result *bizdev.TaskRunResult) {
	ar.mu.Lock()
	ar.record.Status = status
	ar.record.WorkspaceID = workspaceID
	ar.record.ErrorMessage = errMsg
	ar.record.UpdatedAt = time.Now()
	runID := ar.record.RunID
	ar.mu.Unlock()
	slog.Info("taskrun_status_transition",
		"run_id", runID,
		"status", status,
		"workspace_id", workspaceID,
		"error", errMsg,
	)

	// 失败/取消路径：若 Finalize 未执行，仍更新 workflow_runs 与占位消息 metadata
	if status == "failed" || status == "cancelled" || status == "finalizing_failed" {
		if m.pwSvc != nil && ar.record.ProjectID > 0 {
			_ = m.pwSvc.UpdateWorkflowRun(context.Background(), ar.record.RunID, status, "", "", errMsg)
		}
		if m.ctxMgr != nil && ar.record.AssistantMsgID > 0 {
			meta := map[string]any{
				"run_id":       ar.record.RunID,
				"workspace_id": workspaceID,
				"task_status":  status,
			}
			content := "任务执行失败。"
			if status == "cancelled" {
				content = "任务已取消。"
			} else if status == "finalizing_failed" {
				content = "任务已生成，但结果落库失败，请联系管理员或稍后重试。"
			}
			if errMsg != "" {
				content += "\n\n" + errMsg
			}
			_ = m.ctxMgr.FinalizeAssistantTaskMessage(context.Background(), ar.record.AssistantMsgID, content, meta, status)
		}
	}
	_ = result // success path persistence handled by Finalizer
}

func (m *Manager) appendEvent(ar *activeRun, ev stream.StreamEvent) {
	ar.mu.Lock()
	ar.nextSeq++
	stored := StoredEvent{Seq: ar.nextSeq, Event: ev}
	ar.events = append(ar.events, stored)
	subs := make([]chan StoredEvent, 0, len(ar.subs))
	for _, ch := range ar.subs {
		subs = append(subs, ch)
	}
	ar.mu.Unlock()

	for _, ch := range subs {
		select {
		case ch <- stored:
		default:
		}
	}
}

func (m *Manager) GetRun(runID string) (*RunRecord, error) {
	m.mu.RLock()
	ar, ok := m.runs[runID]
	m.mu.RUnlock()
	if !ok {
		return nil, fmt.Errorf("run not found")
	}
	ar.mu.RLock()
	defer ar.mu.RUnlock()
	cpy := *ar.record
	return &cpy, nil
}

func (m *Manager) Subscribe(runID string, afterSeq int64) (<-chan StoredEvent, []StoredEvent, error) {
	m.mu.RLock()
	ar, ok := m.runs[runID]
	m.mu.RUnlock()
	if !ok {
		return nil, nil, fmt.Errorf("run not found")
	}

	ar.mu.Lock()
	defer ar.mu.Unlock()
	replay := make([]StoredEvent, 0)
	for _, ev := range ar.events {
		if ev.Seq > afterSeq {
			replay = append(replay, ev)
		}
	}
	ch := make(chan StoredEvent, 32)
	ar.subSeq++
	id := ar.subSeq
	ar.subs[id] = ch
	return ch, replay, nil
}

func (m *Manager) Unsubscribe(runID string, ch <-chan StoredEvent) {
	m.mu.RLock()
	ar, ok := m.runs[runID]
	m.mu.RUnlock()
	if !ok {
		return
	}
	ar.mu.Lock()
	defer ar.mu.Unlock()
	for id, sub := range ar.subs {
		if sub == ch {
			delete(ar.subs, id)
			close(sub)
			break
		}
	}
}

func (m *Manager) Cancel(runID string) error {
	m.mu.RLock()
	ar, ok := m.runs[runID]
	m.mu.RUnlock()
	if !ok {
		return fmt.Errorf("run not found")
	}
	ar.cancel()
	return nil
}