agentic-coding-workflow/internal/orchestrator/repository.go

package orchestrator

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

	"github.com/google/uuid"
	"github.com/jackc/pgerrcode"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgconn"
	"github.com/jackc/pgx/v5/pgtype"
	"github.com/jackc/pgx/v5/pgxpool"

	"github.com/yan1h/agent-coding-workflow/internal/infra/errs"
	orchsqlc "github.com/yan1h/agent-coding-workflow/internal/orchestrator/sqlc"
	"github.com/yan1h/agent-coding-workflow/internal/project"
)

// Repository 是 orchestrator 模块对 PG 的全部依赖。Service 单测通过手写 fakeRepo 满足。
type Repository interface {
	// Run
	InsertRun(ctx context.Context, r *Run) (*Run, error)
	GetRun(ctx context.Context, id uuid.UUID) (*Run, error)
	GetActiveRunByRequirement(ctx context.Context, requirementID uuid.UUID) (*Run, error)
	ListRuns(ctx context.Context, f RunFilter) ([]*Run, error)
	UpdateRunStatus(ctx context.Context, id uuid.UUID, status RunStatus, lastErr *string) (*Run, error)
	UpdateRunPhase(ctx context.Context, id uuid.UUID, phase project.Phase) (*Run, error)

	// Step
	InsertStep(ctx context.Context, s *Step) (*Step, error)
	GetStep(ctx context.Context, id uuid.UUID) (*Step, error)
	GetStepBySession(ctx context.Context, sessionID uuid.UUID) (*Step, error)
	ListStepsByRun(ctx context.Context, runID uuid.UUID) ([]*Step, error)
	UpdateStepStatus(ctx context.Context, s *Step) (*Step, error)
	IncrementStepAttempt(ctx context.Context, id uuid.UUID) (*Step, error)
	CountActiveChildSteps(ctx context.Context, parentStepID uuid.UUID) (int64, error)
	ResetStuckStepsOnRestart(ctx context.Context) ([]*Step, error)

	InTx(ctx context.Context, fn func(ctx context.Context, tx pgx.Tx) error) error
	WithTx(tx pgx.Tx) Repository
}

// IsUniqueViolation 报告 err 是否为 PG 唯一约束冲突（用于一 requirement 一活跃 run 的 409 翻译）。
func IsUniqueViolation(err error) bool {
	var pg *pgconn.PgError
	return errors.As(err, &pg) && pg.Code == pgerrcode.UniqueViolation
}

type pgRepo struct {
	pool *pgxpool.Pool
	q    *orchsqlc.Queries
}

// NewPostgresRepository 用 pgxpool 构造 Repository。
func NewPostgresRepository(pool *pgxpool.Pool) Repository {
	return &pgRepo{pool: pool, q: orchsqlc.New(pool)}
}

// InTx 在单事务内执行 fn。nil error commit，否则 rollback。StartRun 用它原子地
// 插入 run + 第一个 step。
func (r *pgRepo) InTx(ctx context.Context, fn func(ctx context.Context, tx pgx.Tx) error) error {
	if r.pool == nil {
		return errs.New(errs.CodeInternal, "orchestrator.Repository.InTx: nil pool (tx-bound instance)")
	}
	return pgx.BeginFunc(ctx, r.pool, func(tx pgx.Tx) error {
		return fn(ctx, tx)
	})
}

// WithTx 返回绑定到指定事务的 Repository 副本。
func (r *pgRepo) WithTx(tx pgx.Tx) Repository {
	return &pgRepo{pool: nil, q: r.q.WithTx(tx)}
}

// ===== 类型转换 helper =====

func toPgUUID(u uuid.UUID) pgtype.UUID { return pgtype.UUID{Bytes: u, Valid: true} }
func toPgUUIDPtr(u *uuid.UUID) pgtype.UUID {
	if u == nil {
		return pgtype.UUID{}
	}
	return pgtype.UUID{Bytes: *u, Valid: true}
}
func fromPgUUID(p pgtype.UUID) uuid.UUID {
	if !p.Valid {
		return uuid.Nil
	}
	return uuid.UUID(p.Bytes)
}
func fromPgUUIDPtr(p pgtype.UUID) *uuid.UUID {
	if !p.Valid {
		return nil
	}
	u := uuid.UUID(p.Bytes)
	return &u
}

func timePtr(p pgtype.Timestamptz) *time.Time {
	if !p.Valid {
		return nil
	}
	t := p.Time
	return &t
}

func marshalConfig(c RunConfig) ([]byte, error) {
	b, err := json.Marshal(c)
	if err != nil {
		return nil, errs.Wrap(err, errs.CodeInternal, "marshal run config")
	}
	return b, nil
}

func unmarshalConfig(b []byte) RunConfig {
	var c RunConfig
	if len(b) > 0 {
		_ = json.Unmarshal(b, &c)
	}
	return c
}

func marshalGate(g *GateResult) ([]byte, error) {
	if g == nil {
		return nil, nil
	}
	b, err := json.Marshal(g)
	if err != nil {
		return nil, errs.Wrap(err, errs.CodeInternal, "marshal gate result")
	}
	return b, nil
}

func unmarshalGate(b []byte) *GateResult {
	if len(b) == 0 {
		return nil
	}
	var g GateResult
	if err := json.Unmarshal(b, &g); err != nil {
		return nil
	}
	return &g
}

// ===== Run =====

func (r *pgRepo) InsertRun(ctx context.Context, run *Run) (*Run, error) {
	cfg, err := marshalConfig(run.Config)
	if err != nil {
		return nil, err
	}
	row, err := r.q.InsertRun(ctx, orchsqlc.InsertRunParams{
		ID:            toPgUUID(run.ID),
		ProjectID:     toPgUUID(run.ProjectID),
		RequirementID: toPgUUID(run.RequirementID),
		WorkspaceID:   toPgUUID(run.WorkspaceID),
		OwnerID:       toPgUUID(run.OwnerID),
		Status:        string(run.Status),
		CurrentPhase:  string(run.CurrentPhase),
		Config:        cfg,
		LastError:     run.LastError,
	})
	if err != nil {
		return nil, errs.Wrap(err, errs.CodeInternal, "insert run")
	}
	return rowToRun(row), nil
}

func (r *pgRepo) GetRun(ctx context.Context, id uuid.UUID) (*Run, error) {
	row, err := r.q.GetRun(ctx, toPgUUID(id))
	if err != nil {
		return nil, pgxNoRowsToNotFound(err, "orchestrator run not found")
	}
	return rowToRun(row), nil
}

func (r *pgRepo) GetActiveRunByRequirement(ctx context.Context, requirementID uuid.UUID) (*Run, error) {
	row, err := r.q.GetActiveRunByRequirement(ctx, toPgUUID(requirementID))
	if err != nil {
		return nil, pgxNoRowsToNotFound(err, "no active orchestrator run for requirement")
	}
	return rowToRun(row), nil
}

func (r *pgRepo) ListRuns(ctx context.Context, f RunFilter) ([]*Run, error) {
	limit := int32(f.Limit)
	if limit <= 0 {
		limit = 100
	}
	rows, err := r.q.ListRuns(ctx, orchsqlc.ListRunsParams{
		Column1: toPgUUIDPtr(f.ProjectID),
		Column2: toPgUUIDPtr(f.RequirementID),
		Column3: string(f.Status),
		Limit:   limit,
	})
	if err != nil {
		return nil, errs.Wrap(err, errs.CodeInternal, "list runs")
	}
	out := make([]*Run, 0, len(rows))
	for _, row := range rows {
		out = append(out, rowToRun(row))
	}
	return out, nil
}

func (r *pgRepo) UpdateRunStatus(ctx context.Context, id uuid.UUID, status RunStatus, lastErr *string) (*Run, error) {
	row, err := r.q.UpdateRunStatus(ctx, orchsqlc.UpdateRunStatusParams{
		ID:        toPgUUID(id),
		Status:    string(status),
		LastError: lastErr,
	})
	if err != nil {
		return nil, pgxNoRowsToNotFound(err, "orchestrator run not found")
	}
	return rowToRun(row), nil
}

func (r *pgRepo) UpdateRunPhase(ctx context.Context, id uuid.UUID, phase project.Phase) (*Run, error) {
	row, err := r.q.UpdateRunPhase(ctx, orchsqlc.UpdateRunPhaseParams{
		ID:           toPgUUID(id),
		CurrentPhase: string(phase),
	})
	if err != nil {
		return nil, pgxNoRowsToNotFound(err, "orchestrator run not found")
	}
	return rowToRun(row), nil
}

// ===== Step =====

func (r *pgRepo) InsertStep(ctx context.Context, s *Step) (*Step, error) {
	gate, err := marshalGate(s.GateResult)
	if err != nil {
		return nil, err
	}
	row, err := r.q.InsertStep(ctx, orchsqlc.InsertStepParams{
		ID:           toPgUUID(s.ID),
		RunID:        toPgUUID(s.RunID),
		Phase:        string(s.Phase),
		Attempt:      int32(s.Attempt),
		ParentStepID: toPgUUIDPtr(s.ParentStepID),
		Depth:        int32(s.Depth),
		Status:       string(s.Status),
		AgentKindID:  toPgUUID(s.AgentKindID),
		AcpSessionID: toPgUUIDPtr(s.ACPSessionID),
		Prompt:       s.Prompt,
		StopReason:   s.StopReason,
		GateResult:   gate,
		LastError:    s.LastError,
		JobID:        toPgUUIDPtr(s.JobID),
	})
	if err != nil {
		return nil, errs.Wrap(err, errs.CodeInternal, "insert step")
	}
	return rowToStep(row), nil
}

func (r *pgRepo) GetStep(ctx context.Context, id uuid.UUID) (*Step, error) {
	row, err := r.q.GetStep(ctx, toPgUUID(id))
	if err != nil {
		return nil, pgxNoRowsToNotFound(err, "orchestrator step not found")
	}
	return rowToStep(row), nil
}

func (r *pgRepo) GetStepBySession(ctx context.Context, sessionID uuid.UUID) (*Step, error) {
	row, err := r.q.GetStepBySession(ctx, toPgUUID(sessionID))
	if err != nil {
		return nil, pgxNoRowsToNotFound(err, "orchestrator step not found for session")
	}
	return rowToStep(row), nil
}

func (r *pgRepo) ListStepsByRun(ctx context.Context, runID uuid.UUID) ([]*Step, error) {
	rows, err := r.q.ListStepsByRun(ctx, toPgUUID(runID))
	if err != nil {
		return nil, errs.Wrap(err, errs.CodeInternal, "list steps by run")
	}
	out := make([]*Step, 0, len(rows))
	for _, row := range rows {
		out = append(out, rowToStep(row))
	}
	return out, nil
}

func (r *pgRepo) UpdateStepStatus(ctx context.Context, s *Step) (*Step, error) {
	gate, err := marshalGate(s.GateResult)
	if err != nil {
		return nil, err
	}
	row, err := r.q.UpdateStepStatus(ctx, orchsqlc.UpdateStepStatusParams{
		ID:           toPgUUID(s.ID),
		Status:       string(s.Status),
		AcpSessionID: toPgUUIDPtr(s.ACPSessionID),
		StopReason:   s.StopReason,
		GateResult:   gate,
		LastError:    s.LastError,
		JobID:        toPgUUIDPtr(s.JobID),
	})
	if err != nil {
		return nil, pgxNoRowsToNotFound(err, "orchestrator step not found")
	}
	return rowToStep(row), nil
}

func (r *pgRepo) IncrementStepAttempt(ctx context.Context, id uuid.UUID) (*Step, error) {
	row, err := r.q.IncrementStepAttempt(ctx, toPgUUID(id))
	if err != nil {
		return nil, pgxNoRowsToNotFound(err, "orchestrator step not found")
	}
	return rowToStep(row), nil
}

func (r *pgRepo) CountActiveChildSteps(ctx context.Context, parentStepID uuid.UUID) (int64, error) {
	n, err := r.q.CountActiveChildSteps(ctx, toPgUUIDPtr(&parentStepID))
	if err != nil {
		return 0, errs.Wrap(err, errs.CodeInternal, "count active child steps")
	}
	return n, nil
}

func (r *pgRepo) ResetStuckStepsOnRestart(ctx context.Context) ([]*Step, error) {
	rows, err := r.q.ResetStuckStepsOnRestart(ctx)
	if err != nil {
		return nil, errs.Wrap(err, errs.CodeInternal, "reset stuck steps on restart")
	}
	out := make([]*Step, 0, len(rows))
	for _, row := range rows {
		out = append(out, rowToStep(row))
	}
	return out, nil
}

// ===== row → domain =====

func rowToRun(row orchsqlc.OrchestratorRun) *Run {
	return &Run{
		ID:            fromPgUUID(row.ID),
		ProjectID:     fromPgUUID(row.ProjectID),
		RequirementID: fromPgUUID(row.RequirementID),
		WorkspaceID:   fromPgUUID(row.WorkspaceID),
		OwnerID:       fromPgUUID(row.OwnerID),
		Status:        RunStatus(row.Status),
		CurrentPhase:  project.Phase(row.CurrentPhase),
		Config:        unmarshalConfig(row.Config),
		LastError:     row.LastError,
		CreatedAt:     row.CreatedAt.Time,
		UpdatedAt:     row.UpdatedAt.Time,
		EndedAt:       timePtr(row.EndedAt),
	}
}

func rowToStep(row orchsqlc.OrchestratorStep) *Step {
	return &Step{
		ID:           fromPgUUID(row.ID),
		RunID:        fromPgUUID(row.RunID),
		Phase:        project.Phase(row.Phase),
		Attempt:      int(row.Attempt),
		ParentStepID: fromPgUUIDPtr(row.ParentStepID),
		Depth:        int(row.Depth),
		Status:       StepStatus(row.Status),
		AgentKindID:  fromPgUUID(row.AgentKindID),
		ACPSessionID: fromPgUUIDPtr(row.AcpSessionID),
		Prompt:       row.Prompt,
		StopReason:   row.StopReason,
		GateResult:   unmarshalGate(row.GateResult),
		LastError:    row.LastError,
		JobID:        fromPgUUIDPtr(row.JobID),
		CreatedAt:    row.CreatedAt.Time,
		UpdatedAt:    row.UpdatedAt.Time,
		EndedAt:      timePtr(row.EndedAt),
	}
}

func pgxNoRowsToNotFound(err error, msg string) error {
	if errors.Is(err, pgx.ErrNoRows) {
		return errs.New(errs.CodeNotFound, msg)
	}
	return errs.Wrap(err, errs.CodeInternal, msg)
}