Agent-System/internal/taskrun/manager.go

package taskrun

import (
	"context"
	"errors"
	"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；可选 PG 持久化支持重启重放。
type Manager struct {
	mu       sync.RWMutex
	runs     map[string]*activeRun
	store    Store
	workerID string
	exec     Executor
	resume   ResumeExecutor
	stage    StageExecutor
	finalize Finalizer
	cmds     *Commands
	pwSvc    *projectworkflow.Service
	ctxMgr   *ctxmgr.ContextManager

	idemMu    sync.Mutex
	idemLocks map[string]*sync.Mutex // keyed by idempotency_key
}

const runRetention = time.Hour

var (
	ErrRunNotResumable  = errors.New("taskrun: run is not resumable")
	ErrRunAlreadyActive = errors.New("taskrun: run already active in memory")
)

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

func (m *Manager) lockIdempotencyKey(key string) func() {
	if key == "" {
		return func() {}
	}
	m.idemMu.Lock()
	lk, ok := m.idemLocks[key]
	if !ok {
		lk = &sync.Mutex{}
		m.idemLocks[key] = lk
	}
	m.idemMu.Unlock()
	lk.Lock()
	return lk.Unlock
}

// SetStore 注入持久化存储（PG）；启动时标记 lease 过期的 running 任务为 interrupted。
func (m *Manager) SetStore(store Store) {
	if m == nil {
		return
	}
	if store == nil {
		store = noopStore{}
	}
	m.store = store
	if store.Available() {
		if n, err := store.MarkOrphanedRunning(context.Background()); err != nil {
			slog.Warn("taskrun_orphan_recovery_failed", "error", err)
		} else if n > 0 {
			slog.Info("taskrun_orphans_marked_interrupted", "count", n)
		}
	}
}

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) SetCommands(cmds *Commands) {
	m.cmds = cmds
	if cmds != nil {
		m.finalize = cmds.Finalize
	}
}

func (m *Manager) CreateAndStart(parent context.Context, runID string, record *RunRecord, taskReq bizdev.TaskRunRequest, stageReq bizdev.StageRunRequest, resumeWorkDir string) (*RunRecord, error) {
	if record == nil {
		record = &RunRecord{}
	}
	if runID != "" {
		record.RunID = runID
	}
	claimed, created, err := m.ClaimRun(parent, record)
	if err != nil {
		return nil, err
	}
	if !created {
		return claimed, nil
	}
	return m.StartClaimedRun(parent, claimed, taskReq, stageReq, resumeWorkDir)
}

// ClaimRun 原子占位：在写聊天 / 启 worker 之前确认本请求是否为 owner。
// created=false 表示已有同 idempotency_key 的 run，调用方不得再写消息。
func (m *Manager) ClaimRun(parent context.Context, record *RunRecord) (*RunRecord, bool, error) {
	if record == nil {
		return nil, false, fmt.Errorf("taskrun: record required")
	}
	if record.IdempotencyKey == "" && record.SessionID != "" && record.ClientTurnID != "" {
		record.IdempotencyKey = BuildTaskRunIdempotencyKey(record.SessionID, record.ClientTurnID)
	}
	unlock := m.lockIdempotencyKey(record.IdempotencyKey)
	defer unlock()

	if record.IdempotencyKey != "" && m.store != nil && m.store.Available() {
		if existing, err := m.store.GetByIdempotencyKey(parent, record.IdempotencyKey); err != nil {
			return nil, false, err
		} else if existing != nil {
			return existing, false, nil
		}
	}

	if record.RunID == "" {
		record.RunID = fmt.Sprintf("run_%d", time.Now().UnixNano())
	}
	if record.Status == "" {
		record.Status = "pending"
	}
	now := time.Now()
	if record.CreatedAt.IsZero() {
		record.CreatedAt = now
	}
	record.UpdatedAt = now
	if record.Attempt <= 0 {
		record.Attempt = 1
	}

	if m.store != nil && m.store.Available() {
		if err := m.store.InsertRun(parent, record); err != nil {
			if errors.Is(err, ErrIdempotencyConflict) && record.IdempotencyKey != "" {
				existing, lookupErr := m.store.GetByIdempotencyKey(parent, record.IdempotencyKey)
				if lookupErr != nil {
					return nil, false, lookupErr
				}
				if existing != nil {
					return existing, false, nil
				}
			}
			return nil, false, err
		}
	}
	cpy := *record
	return &cpy, true, nil
}

// BindAssistantMsgID 在启 worker 前回填 assistant 消息 id，覆盖快速完成竞态。
func (m *Manager) BindAssistantMsgID(ctx context.Context, runID string, assistantMsgID int64) error {
	if runID == "" || assistantMsgID <= 0 {
		return nil
	}
	m.mu.RLock()
	ar, ok := m.runs[runID]
	m.mu.RUnlock()
	if ok {
		ar.mu.Lock()
		ar.record.AssistantMsgID = assistantMsgID
		ar.record.UpdatedAt = time.Now()
		rec := *ar.record
		ar.mu.Unlock()
		if m.store != nil && m.store.Available() {
			return m.store.UpdateRun(ctx, &rec)
		}
		return nil
	}
	if m.store != nil && m.store.Available() {
		rec, err := m.store.GetRun(ctx, runID)
		if err != nil {
			return err
		}
		rec.AssistantMsgID = assistantMsgID
		rec.UpdatedAt = time.Now()
		return m.store.UpdateRun(ctx, rec)
	}
	return nil
}

// StartClaimedRun 仅对 ClaimRun 成功的 owner 启 worker；禁止在 Insert 冲突后调用。
func (m *Manager) StartClaimedRun(parent context.Context, record *RunRecord, taskReq bizdev.TaskRunRequest, stageReq bizdev.StageRunRequest, resumeWorkDir string) (*RunRecord, error) {
	if record == nil || record.RunID == "" {
		return nil, fmt.Errorf("taskrun: claimed record required")
	}
	record.Status = "running"
	record.WorkerID = m.workerID
	now := time.Now()
	record.UpdatedAt = now
	leaseUntil := now.Add(defaultLeaseTTL)
	record.LeaseExpiresAt = &leaseUntil

	if m.store != nil && m.store.Available() {
		if err := m.store.UpdateRun(parent, record); err != nil {
			return nil, fmt.Errorf("taskrun: persist running status before start: %w", err)
		}
		if err := m.store.AcquireLease(parent, record.RunID, m.workerID, defaultLeaseTTL); err != nil {
			slog.Warn("taskrun_lease_acquire_failed", "run_id", record.RunID, "error", err)
		}
	}

	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[record.RunID] = ar
	m.mu.Unlock()

	if m.pwSvc != nil && record.ProjectID > 0 {
		startedBy := record.UserID
		_, _ = m.pwSvc.CreateWorkflowRun(workerCtx, &projectworkflow.WorkflowRun{
			RunID:     record.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)
	cpy := *record
	return &cpy, nil
}

func (m *Manager) runWorker(ctx context.Context, ar *activeRun, taskReq bizdev.TaskRunRequest, stageReq bizdev.StageRunRequest, resumeWorkDir string) {
	defer close(ar.done)
	hbCtx, hbCancel := context.WithCancel(ctx)
	defer hbCancel()
	if m.store != nil && m.store.Available() {
		go m.heartbeatLoop(hbCtx, ar.record.RunID)
	}

	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.record.UpdatedAt = time.Now()
		ar.mu.Unlock()
		m.persistRun(ar)
	}

	if execErr != nil {
		if ctx.Err() == context.Canceled {
			if m.cmds != nil {
				_, _ = m.cmds.Cancel(context.Background(), ar.record)
			}
			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})
			m.finishRun(ar, "cancelled", workspaceID, execErr.Error(), nil, &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})
		if m.cmds != nil {
			failRes := taskResult
			if failRes == nil {
				failRes = &bizdev.TaskRunResult{Success: false, ErrorReason: execErr.Error(), WorkDir: workspaceID}
			} else {
				failRes.Success = false
				if failRes.ErrorReason == "" {
					failRes.ErrorReason = execErr.Error()
				}
			}
			_, _ = m.cmds.Finalize(context.Background(), ar.record, failRes)
		}
		m.finishRun(ar, "failed", workspaceID, execErr.Error(), taskResult, &info)
		return
	}

	// RUN-002: 先 finalize 落库，成功后再发 done 终态事件
	terminalStatus := "completed"
	if m.finalize != nil && taskResult != nil {
		if ts, finErr := m.finalize(ctx, ar.record, taskResult); finErr != nil {
			info := stream.ClassifyError(finErr, lastPhaseID, ar.record.RunID, ar.record.SessionID, workspaceID)
			m.appendEvent(ar, stream.StreamEvent{Type: stream.EventError, ErrorInfo: &info})
			m.finishRun(ar, "finalizing_failed", workspaceID, finErr.Error(), taskResult, &info)
			return
		} else if ts != "" {
			terminalStatus = ts
		}
	} else if taskResult != nil && taskResult.QualityReport != nil && !taskResult.QualityReport.Passed {
		terminalStatus = "quality_failed"
	}

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

func (m *Manager) finishRun(ar *activeRun, status, workspaceID, errMsg string, result *bizdev.TaskRunResult, errorInfo *stream.ErrorInfo) {
	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,
	)
	if m.store != nil && m.store.Available() {
		_ = m.store.ReleaseLease(context.Background(), runID, m.workerID)
	}
	ar.mu.Lock()
	ar.record.WorkerID = ""
	ar.record.LeaseExpiresAt = nil
	ar.mu.Unlock()
	m.persistRun(ar)

	// 失败/取消路径：若 Finalize 未执行，仍更新 workflow_runs 与占位消息 metadata
	if m.cmds == nil && (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,
			}
			if errorInfo != nil {
				meta["error_info"] = errorInfo
			}
			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) persistRun(ar *activeRun) {
	if m.store == nil || !m.store.Available() || ar == nil {
		return
	}
	ar.mu.RLock()
	rec := *ar.record
	ar.mu.RUnlock()
	if err := m.store.UpdateRun(context.Background(), &rec); err != nil {
		slog.Warn("taskrun_update_failed", "run_id", rec.RunID, "error", err)
	}
}

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)
	runID := ar.record.RunID
	seq := ar.nextSeq
	subs := make([]chan StoredEvent, 0, len(ar.subs))
	for _, ch := range ar.subs {
		subs = append(subs, ch)
	}
	ar.mu.Unlock()

	if m.store != nil && m.store.Available() {
		if err := m.store.AppendEvent(context.Background(), runID, seq, ev); err != nil {
			slog.Warn("taskrun_event_persist_failed", "run_id", runID, "seq", seq, "error", err)
		}
	}

	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 {
		ar.mu.RLock()
		defer ar.mu.RUnlock()
		cpy := *ar.record
		return &cpy, nil
	}
	if m.store != nil && m.store.Available() {
		return m.store.GetRun(context.Background(), runID)
	}
	return nil, fmt.Errorf("run not found")
}

// FindByIdempotencyKey 查找既有幂等 run（内存优先，再查 store）。
func (m *Manager) FindByIdempotencyKey(ctx context.Context, key string) (*RunRecord, error) {
	if key == "" {
		return nil, nil
	}
	m.mu.RLock()
	for _, ar := range m.runs {
		ar.mu.RLock()
		match := ar.record != nil && ar.record.IdempotencyKey == key
		var cpy RunRecord
		if match {
			cpy = *ar.record
		}
		ar.mu.RUnlock()
		if match {
			m.mu.RUnlock()
			return &cpy, nil
		}
	}
	m.mu.RUnlock()
	if m.store != nil && m.store.Available() {
		return m.store.GetByIdempotencyKey(ctx, key)
	}
	return nil, 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 {
		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
	}

	if m.store != nil && m.store.Available() {
		if _, err := m.store.GetRun(context.Background(), runID); err != nil {
			return nil, nil, fmt.Errorf("run not found")
		}
		events, err := m.store.ListEvents(context.Background(), runID, afterSeq)
		if err != nil {
			return nil, nil, err
		}
		ch := make(chan StoredEvent)
		close(ch)
		return ch, events, nil
	}
	return nil, nil, fmt.Errorf("run not found")
}

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

// ResumeInterrupted 在用户确认后续跑 interrupted 任务（需有效 workspace 路径）。
func (m *Manager) ResumeInterrupted(parent context.Context, record *RunRecord, resumeWorkDir string, taskReq bizdev.TaskRunRequest, stageReq bizdev.StageRunRequest) (*RunRecord, error) {
	if record == nil || record.RunID == "" {
		return nil, fmt.Errorf("taskrun: run_id required")
	}
	if resumeWorkDir == "" {
		return nil, fmt.Errorf("taskrun: workspace required")
	}
	if record.Status != "interrupted" {
		return nil, ErrRunNotResumable
	}

	m.mu.RLock()
	_, exists := m.runs[record.RunID]
	m.mu.RUnlock()
	if exists {
		return nil, ErrRunAlreadyActive
	}

	nextAttempt := record.Attempt + 1
	if nextAttempt <= 0 {
		nextAttempt = 2
	}
	if m.store != nil && m.store.Available() {
		if err := m.store.ResetForResume(parent, record.RunID, nextAttempt); err != nil {
			return nil, err
		}
	}

	now := time.Now()
	leaseUntil := now.Add(defaultLeaseTTL)
	cpy := *record
	cpy.Status = "running"
	cpy.ErrorMessage = ""
	cpy.Attempt = nextAttempt
	cpy.WorkerID = m.workerID
	cpy.LeaseExpiresAt = &leaseUntil
	cpy.UpdatedAt = now

	if m.store != nil && m.store.Available() {
		_ = m.store.UpdateRun(parent, &cpy)
		if err := m.store.AcquireLease(parent, cpy.RunID, m.workerID, defaultLeaseTTL); err != nil {
			slog.Warn("taskrun_resume_lease_failed", "run_id", cpy.RunID, "error", err)
		}
	}

	workerCtx, cancel := context.WithCancel(context.WithoutCancel(parent))
	ar := &activeRun{
		record: &cpy,
		subs:   make(map[int]chan StoredEvent),
		cancel: cancel,
		done:   make(chan struct{}),
	}

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

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

func (m *Manager) Cancel(runID string) error {
	rec, _ := m.GetRun(runID)
	if m.cmds != nil && rec != nil {
		if _, err := m.cmds.Cancel(context.Background(), rec); err != nil && !errors.Is(err, ErrTerminalConflict) {
			return err
		}
	}
	m.mu.RLock()
	ar, ok := m.runs[runID]
	m.mu.RUnlock()
	if ok {
		ar.cancel()
		return nil
	}
	if rec != nil {
		return nil
	}
	return fmt.Errorf("run not found")
}

func (m *Manager) heartbeatLoop(ctx context.Context, runID string) {
	ticker := time.NewTicker(defaultHeartbeatPeriod)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			if err := m.store.HeartbeatLease(context.Background(), runID, m.workerID, defaultLeaseTTL); err != nil {
				slog.Warn("taskrun_heartbeat_failed", "run_id", runID, "error", err)
			}
		}
	}
}