Agent-System/internal/taskrun/store.go

package taskrun

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"

	"chat-agent-system/internal/stream"
)

var ErrStoreNotAvailable = errors.New("taskrun: store not available")

// Store 持久化 task run 与 SSE 事件，支持重启后查询与重放。
type Store interface {
	Available() bool
	InsertRun(ctx context.Context, rec *RunRecord) error
	UpdateRun(ctx context.Context, rec *RunRecord) error
	AppendEvent(ctx context.Context, runID string, seq int64, ev stream.StreamEvent) error
	GetRun(ctx context.Context, runID string) (*RunRecord, error)
	ListEvents(ctx context.Context, runID string, afterSeq int64) ([]StoredEvent, error)
	AcquireLease(ctx context.Context, runID, workerID string, ttl time.Duration) error
	HeartbeatLease(ctx context.Context, runID, workerID string, ttl time.Duration) error
	ReleaseLease(ctx context.Context, runID, workerID string) error
	ResetForResume(ctx context.Context, runID string, attempt int) error
	MarkOrphanedRunning(ctx context.Context) (int, error)
}

// PGStore PostgreSQL 实现。
type PGStore struct {
	pool *pgxpool.Pool
}

func NewPGStore(pool *pgxpool.Pool) *PGStore {
	return &PGStore{pool: pool}
}

func (s *PGStore) Available() bool {
	if s == nil || s.pool == nil {
		return false
	}
	var ok bool
	_ = s.pool.QueryRow(context.Background(), `
		SELECT EXISTS (
			SELECT 1 FROM information_schema.tables
			WHERE table_schema = 'public' AND table_name = 'task_runs'
		)
	`).Scan(&ok)
	return ok
}

func (s *PGStore) hasLeaseColumns(ctx context.Context) bool {
	var ok bool
	_ = s.pool.QueryRow(ctx, `
		SELECT EXISTS (
			SELECT 1 FROM information_schema.columns
			WHERE table_schema = 'public' AND table_name = 'task_runs' AND column_name = 'worker_id'
		)
	`).Scan(&ok)
	return ok
}

func (s *PGStore) InsertRun(ctx context.Context, rec *RunRecord) error {
	if !s.Available() {
		return ErrStoreNotAvailable
	}
	if rec == nil || rec.RunID == "" {
		return fmt.Errorf("taskrun: run_id required")
	}
	attempt := rec.Attempt
	if attempt <= 0 {
		attempt = 1
	}
	if s.hasLeaseColumns(ctx) {
		var leaseHeartbeat *time.Time
		if rec.LeaseExpiresAt != nil {
			now := time.Now()
			leaseHeartbeat = &now
		}
		_, err := s.pool.Exec(ctx, `
			INSERT INTO task_runs (
				run_id, session_id, project_id, user_id, client_turn_id,
				task_key, stage, status, workspace_id, assistant_message_id,
				error_message, worker_id, lease_expires_at, lease_heartbeat_at, attempt,
				created_at, updated_at
			) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17)
			ON CONFLICT (run_id) DO NOTHING
		`,
			rec.RunID, rec.SessionID, nullInt64(rec.ProjectID), rec.UserID, rec.ClientTurnID,
			rec.TaskKey, rec.Stage, rec.Status, rec.WorkspaceID, nullInt64(rec.AssistantMsgID),
			rec.ErrorMessage, rec.WorkerID, rec.LeaseExpiresAt, leaseHeartbeat, attempt,
			rec.CreatedAt, rec.UpdatedAt,
		)
		return err
	}
	_, err := s.pool.Exec(ctx, `
		INSERT INTO task_runs (
			run_id, session_id, project_id, user_id, client_turn_id,
			task_key, stage, status, workspace_id, assistant_message_id,
			error_message, created_at, updated_at
		) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13)
		ON CONFLICT (run_id) DO NOTHING
	`,
		rec.RunID, rec.SessionID, nullInt64(rec.ProjectID), rec.UserID, rec.ClientTurnID,
		rec.TaskKey, rec.Stage, rec.Status, rec.WorkspaceID, nullInt64(rec.AssistantMsgID),
		rec.ErrorMessage, rec.CreatedAt, rec.UpdatedAt,
	)
	return err
}

func (s *PGStore) UpdateRun(ctx context.Context, rec *RunRecord) error {
	if !s.Available() {
		return ErrStoreNotAvailable
	}
	if rec == nil || rec.RunID == "" {
		return fmt.Errorf("taskrun: run_id required")
	}
	if s.hasLeaseColumns(ctx) {
		_, err := s.pool.Exec(ctx, `
			UPDATE task_runs SET
				status = $2,
				workspace_id = $3,
				error_message = $4,
				worker_id = $5,
				lease_expires_at = $6,
				updated_at = $7
			WHERE run_id = $1
		`, rec.RunID, rec.Status, rec.WorkspaceID, rec.ErrorMessage, rec.WorkerID, rec.LeaseExpiresAt, rec.UpdatedAt)
		return err
	}
	_, err := s.pool.Exec(ctx, `
		UPDATE task_runs SET
			status = $2,
			workspace_id = $3,
			error_message = $4,
			updated_at = $5
		WHERE run_id = $1
	`, rec.RunID, rec.Status, rec.WorkspaceID, rec.ErrorMessage, rec.UpdatedAt)
	return err
}

func (s *PGStore) AppendEvent(ctx context.Context, runID string, seq int64, ev stream.StreamEvent) error {
	if !s.Available() {
		return ErrStoreNotAvailable
	}
	payload, err := json.Marshal(ev)
	if err != nil {
		return err
	}
	_, err = s.pool.Exec(ctx, `
		INSERT INTO task_run_events (run_id, seq, event_type, payload)
		VALUES ($1, $2, $3, $4::jsonb)
		ON CONFLICT (run_id, seq) DO NOTHING
	`, runID, seq, string(ev.Type), payload)
	return err
}

func (s *PGStore) GetRun(ctx context.Context, runID string) (*RunRecord, error) {
	if !s.Available() {
		return nil, ErrStoreNotAvailable
	}
	var rec RunRecord
	var projectID, assistantMsgID *int64
	if s.hasLeaseColumns(ctx) {
		err := s.pool.QueryRow(ctx, `
			SELECT run_id, session_id, project_id, user_id, client_turn_id,
				task_key, stage, status, workspace_id, assistant_message_id,
				error_message, worker_id, lease_expires_at, attempt, created_at, updated_at
			FROM task_runs
			WHERE run_id = $1
		`, runID).Scan(
			&rec.RunID, &rec.SessionID, &projectID, &rec.UserID, &rec.ClientTurnID,
			&rec.TaskKey, &rec.Stage, &rec.Status, &rec.WorkspaceID, &assistantMsgID,
			&rec.ErrorMessage, &rec.WorkerID, &rec.LeaseExpiresAt, &rec.Attempt,
			&rec.CreatedAt, &rec.UpdatedAt,
		)
		if err != nil {
			if errors.Is(err, pgx.ErrNoRows) {
				return nil, fmt.Errorf("run not found")
			}
			return nil, err
		}
	} else {
		err := s.pool.QueryRow(ctx, `
			SELECT run_id, session_id, project_id, user_id, client_turn_id,
				task_key, stage, status, workspace_id, assistant_message_id,
				error_message, created_at, updated_at
			FROM task_runs
			WHERE run_id = $1
		`, runID).Scan(
			&rec.RunID, &rec.SessionID, &projectID, &rec.UserID, &rec.ClientTurnID,
			&rec.TaskKey, &rec.Stage, &rec.Status, &rec.WorkspaceID, &assistantMsgID,
			&rec.ErrorMessage, &rec.CreatedAt, &rec.UpdatedAt,
		)
		if err != nil {
			if errors.Is(err, pgx.ErrNoRows) {
				return nil, fmt.Errorf("run not found")
			}
			return nil, err
		}
	}
	if projectID != nil {
		rec.ProjectID = *projectID
	}
	if assistantMsgID != nil {
		rec.AssistantMsgID = *assistantMsgID
	}
	return &rec, nil
}

func (s *PGStore) ListEvents(ctx context.Context, runID string, afterSeq int64) ([]StoredEvent, error) {
	if !s.Available() {
		return nil, ErrStoreNotAvailable
	}
	rows, err := s.pool.Query(ctx, `
		SELECT seq, payload
		FROM task_run_events
		WHERE run_id = $1 AND seq > $2
		ORDER BY seq ASC
	`, runID, afterSeq)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	out := make([]StoredEvent, 0)
	for rows.Next() {
		var seq int64
		var payload []byte
		if err := rows.Scan(&seq, &payload); err != nil {
			return nil, err
		}
		var ev stream.StreamEvent
		if err := json.Unmarshal(payload, &ev); err != nil {
			return nil, err
		}
		out = append(out, StoredEvent{Seq: seq, Event: ev})
	}
	return out, rows.Err()
}

func (s *PGStore) AcquireLease(ctx context.Context, runID, workerID string, ttl time.Duration) error {
	if !s.Available() {
		return ErrStoreNotAvailable
	}
	if !s.hasLeaseColumns(ctx) {
		return nil
	}
	expires := time.Now().Add(ttl)
	tag, err := s.pool.Exec(ctx, `
		UPDATE task_runs
		SET worker_id = $2,
		    lease_expires_at = $3,
		    lease_heartbeat_at = now(),
		    updated_at = now()
		WHERE run_id = $1 AND status = 'running'
	`, runID, workerID, expires)
	if err != nil {
		return err
	}
	if tag.RowsAffected() == 0 {
		return fmt.Errorf("taskrun: lease acquire failed for %s", runID)
	}
	return nil
}

func (s *PGStore) HeartbeatLease(ctx context.Context, runID, workerID string, ttl time.Duration) error {
	if !s.Available() {
		return ErrStoreNotAvailable
	}
	if !s.hasLeaseColumns(ctx) {
		return nil
	}
	expires := time.Now().Add(ttl)
	tag, err := s.pool.Exec(ctx, `
		UPDATE task_runs
		SET lease_expires_at = $3,
		    lease_heartbeat_at = now(),
		    updated_at = now()
		WHERE run_id = $1 AND worker_id = $2 AND status = 'running'
	`, runID, workerID, expires)
	if err != nil {
		return err
	}
	if tag.RowsAffected() == 0 {
		return fmt.Errorf("taskrun: heartbeat rejected for %s", runID)
	}
	return nil
}

func (s *PGStore) ReleaseLease(ctx context.Context, runID, workerID string) error {
	if !s.Available() {
		return ErrStoreNotAvailable
	}
	if !s.hasLeaseColumns(ctx) {
		return nil
	}
	_, err := s.pool.Exec(ctx, `
		UPDATE task_runs
		SET worker_id = '',
		    lease_expires_at = NULL,
		    lease_heartbeat_at = NULL,
		    updated_at = now()
		WHERE run_id = $1 AND (worker_id = $2 OR worker_id = '')
	`, runID, workerID)
	return err
}

func (s *PGStore) ResetForResume(ctx context.Context, runID string, attempt int) error {
	if !s.Available() {
		return ErrStoreNotAvailable
	}
	tag, err := s.pool.Exec(ctx, `
		UPDATE task_runs
		SET status = 'running',
		    error_message = '',
		    attempt = $2,
		    updated_at = now()
		WHERE run_id = $1 AND status = 'interrupted'
	`, runID, attempt)
	if err != nil {
		return err
	}
	if tag.RowsAffected() == 0 {
		return fmt.Errorf("taskrun: run %s not resumable", runID)
	}
	return nil
}

func (s *PGStore) MarkOrphanedRunning(ctx context.Context) (int, error) {
	if !s.Available() {
		return 0, ErrStoreNotAvailable
	}
	if s.hasLeaseColumns(ctx) {
		tag, err := s.pool.Exec(ctx, `
			UPDATE task_runs
			SET status = 'interrupted',
			    error_message = 'worker lease 过期，任务已中断（进程重启或 worker 失联）',
			    worker_id = '',
			    lease_expires_at = NULL,
			    lease_heartbeat_at = NULL,
			    updated_at = now()
			WHERE status = 'running'
			  AND lease_expires_at IS NOT NULL
			  AND lease_expires_at < now() - interval '2 seconds'
		`)
		if err != nil {
			return 0, err
		}
		return int(tag.RowsAffected()), nil
	}
	tag, err := s.pool.Exec(ctx, `
		UPDATE task_runs
		SET status = 'interrupted',
		    error_message = '进程重启导致任务中断，请重新提交',
		    updated_at = now()
		WHERE status = 'running'
	`)
	if err != nil {
		return 0, err
	}
	return int(tag.RowsAffected()), nil
}

func nullInt64(v int64) *int64 {
	if v <= 0 {
		return nil
	}
	return &v
}

// TablesReady 检查 task_runs 表是否已迁移。
func TablesReady(ctx context.Context, pool *pgxpool.Pool) bool {
	if pool == nil {
		return false
	}
	var ok bool
	_ = pool.QueryRow(ctx, `
		SELECT EXISTS (
			SELECT 1 FROM information_schema.tables
			WHERE table_schema = 'public' AND table_name = 'task_runs'
		)
	`).Scan(&ok)
	return ok
}

// LeaseColumnsReady 检查 worker lease 列是否已迁移。
func LeaseColumnsReady(ctx context.Context, pool *pgxpool.Pool) bool {
	if pool == nil {
		return false
	}
	var ok bool
	_ = pool.QueryRow(ctx, `
		SELECT EXISTS (
			SELECT 1 FROM information_schema.columns
			WHERE table_schema = 'public' AND table_name = 'task_runs' AND column_name = 'worker_id'
		)
	`).Scan(&ok)
	return ok
}

// noopStore 无持久化时的空实现。
type noopStore struct{}

func (noopStore) Available() bool { return false }
func (noopStore) InsertRun(context.Context, *RunRecord) error {
	return ErrStoreNotAvailable
}
func (noopStore) UpdateRun(context.Context, *RunRecord) error { return ErrStoreNotAvailable }
func (noopStore) AppendEvent(context.Context, string, int64, stream.StreamEvent) error {
	return ErrStoreNotAvailable
}
func (noopStore) GetRun(context.Context, string) (*RunRecord, error) {
	return nil, ErrStoreNotAvailable
}
func (noopStore) ListEvents(context.Context, string, int64) ([]StoredEvent, error) {
	return nil, ErrStoreNotAvailable
}
func (noopStore) AcquireLease(context.Context, string, string, time.Duration) error {
	return ErrStoreNotAvailable
}
func (noopStore) HeartbeatLease(context.Context, string, string, time.Duration) error {
	return ErrStoreNotAvailable
}
func (noopStore) ReleaseLease(context.Context, string, string) error { return ErrStoreNotAvailable }
func (noopStore) ResetForResume(context.Context, string, int) error {
	return ErrStoreNotAvailable
}
func (noopStore) MarkOrphanedRunning(context.Context) (int, error) { return 0, ErrStoreNotAvailable }

// EnsureUpdatedAt 若 UpdatedAt 为零则设为当前时间。
func EnsureUpdatedAt(rec *RunRecord) {
	if rec != nil && rec.UpdatedAt.IsZero() {
		rec.UpdatedAt = time.Now()
	}
}