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

package acp

import (
	"context"
	"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"

	acpsqlc "github.com/yan1h/agent-coding-workflow/internal/acp/sqlc"
	"github.com/yan1h/agent-coding-workflow/internal/infra/errs"
)

// Repository 是 acp 模块对 PG 的全部依赖。Service 单测通过手写 fakeRepo 满足。
type Repository interface {
	// AgentKind
	CreateAgentKind(ctx context.Context, k *AgentKind) (*AgentKind, error)
	GetAgentKindByID(ctx context.Context, id uuid.UUID) (*AgentKind, error)
	GetAgentKindByName(ctx context.Context, name string) (*AgentKind, error)
	ListAgentKinds(ctx context.Context) ([]*AgentKind, error)
	ListEnabledAgentKinds(ctx context.Context) ([]*AgentKind, error)
	UpdateAgentKind(ctx context.Context, k *AgentKind) (*AgentKind, error)
	DeleteAgentKind(ctx context.Context, id uuid.UUID) error
	CountAgentKindUsage(ctx context.Context, id uuid.UUID) (int64, error)

	// Session
	InsertSession(ctx context.Context, s *Session) (*Session, error)
	GetSessionByID(ctx context.Context, id uuid.UUID) (*Session, error)
	ListSessionsByUser(ctx context.Context, userID uuid.UUID) ([]*Session, error)
	ListAllSessions(ctx context.Context) ([]*Session, error)
	UpdateSessionRunning(ctx context.Context, id uuid.UUID, agentSessionID string, pid int32) error
	UpdateSessionFinished(ctx context.Context, id uuid.UUID, status SessionStatus, exitCode *int32, lastErr *string) error
	CountActiveSessions(ctx context.Context) (int64, error)
	CountActiveSessionsByUser(ctx context.Context, userID uuid.UUID) (int64, error)
	ResetStuckSessionsOnRestart(ctx context.Context) ([]*Session, error)

	// Worktree mutex（main worktree 专用）
	AcquireMainWorktree(ctx context.Context, sessionID, workspaceID uuid.UUID) (bool, error)
	ReleaseMainWorktree(ctx context.Context, workspaceID, sessionID uuid.UUID) (bool, error)

	// Event
	InsertEvent(ctx context.Context, e *Event) (*Event, error)
	ListEventsSince(ctx context.Context, sessionID uuid.UUID, sinceID int64, limit int32) ([]*Event, error)
	PurgeEventsBefore(ctx context.Context, before time.Time) (int64, error)

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

// IsUniqueViolation reports whether err is a PG unique constraint violation
// (used by service layer to translate to a 409 if needed).
func IsUniqueViolation(err error) bool {
	var pg *pgconn.PgError
	return errors.As(err, &pg) && pg.Code == pgerrcode.UniqueViolation
}

// IsForeignKeyViolation reports whether err is a PG FK violation. ACP uses
// this when delete is rejected because acp_sessions still references the
// agent_kind (ON DELETE RESTRICT triggers).
func IsForeignKeyViolation(err error) bool {
	var pg *pgconn.PgError
	return errors.As(err, &pg) && pg.Code == pgerrcode.ForeignKeyViolation
}

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

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

// InTx runs fn in a single transaction. Commits on nil error; rolls back
// otherwise. Used by SessionService.Create to atomically insert acp_sessions
// + mcp_tokens (spec §6.1).
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, "acp.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
	}
	id := uuid.UUID(p.Bytes)
	return &id
}

// pgxNoRowsToErrNotFound 把 pgx.ErrNoRows 翻译为 errs.AppError(NotFound)。
func pgxNoRowsToErrNotFound(err error, code errs.Code, msg string) error {
	if errors.Is(err, pgx.ErrNoRows) {
		return errs.New(code, msg)
	}
	return errs.Wrap(err, errs.CodeInternal, msg)
}

// ===== AgentKind =====

func (r *pgRepo) CreateAgentKind(ctx context.Context, k *AgentKind) (*AgentKind, error) {
	row, err := r.q.CreateAgentKind(ctx, acpsqlc.CreateAgentKindParams{
		ID:           toPgUUID(k.ID),
		Name:         k.Name,
		DisplayName:  k.DisplayName,
		Description:  k.Description,
		BinaryPath:   k.BinaryPath,
		Args:         k.Args,
		EncryptedEnv: k.EncryptedEnv,
		Enabled:      k.Enabled,
		CreatedBy:    toPgUUID(k.CreatedBy),
	})
	if err != nil {
		return nil, errs.Wrap(err, errs.CodeInternal, "create agent kind")
	}
	return rowToAgentKind(row), nil
}

func (r *pgRepo) GetAgentKindByID(ctx context.Context, id uuid.UUID) (*AgentKind, error) {
	row, err := r.q.GetAgentKindByID(ctx, toPgUUID(id))
	if err != nil {
		return nil, pgxNoRowsToErrNotFound(err, errs.CodeAcpAgentKindNotFound, "agent kind not found")
	}
	return rowToAgentKind(row), nil
}

func (r *pgRepo) GetAgentKindByName(ctx context.Context, name string) (*AgentKind, error) {
	row, err := r.q.GetAgentKindByName(ctx, name)
	if err != nil {
		return nil, pgxNoRowsToErrNotFound(err, errs.CodeAcpAgentKindNotFound, "agent kind not found")
	}
	return rowToAgentKind(row), nil
}

func (r *pgRepo) ListAgentKinds(ctx context.Context) ([]*AgentKind, error) {
	rows, err := r.q.ListAgentKinds(ctx)
	if err != nil {
		return nil, errs.Wrap(err, errs.CodeInternal, "list agent kinds")
	}
	out := make([]*AgentKind, 0, len(rows))
	for _, row := range rows {
		out = append(out, rowToAgentKind(row))
	}
	return out, nil
}

func (r *pgRepo) ListEnabledAgentKinds(ctx context.Context) ([]*AgentKind, error) {
	rows, err := r.q.ListEnabledAgentKinds(ctx)
	if err != nil {
		return nil, errs.Wrap(err, errs.CodeInternal, "list enabled agent kinds")
	}
	out := make([]*AgentKind, 0, len(rows))
	for _, row := range rows {
		out = append(out, rowToAgentKind(row))
	}
	return out, nil
}

func (r *pgRepo) UpdateAgentKind(ctx context.Context, k *AgentKind) (*AgentKind, error) {
	row, err := r.q.UpdateAgentKind(ctx, acpsqlc.UpdateAgentKindParams{
		ID:           toPgUUID(k.ID),
		DisplayName:  k.DisplayName,
		Description:  k.Description,
		BinaryPath:   k.BinaryPath,
		Args:         k.Args,
		EncryptedEnv: k.EncryptedEnv,
		Enabled:      k.Enabled,
	})
	if err != nil {
		return nil, pgxNoRowsToErrNotFound(err, errs.CodeAcpAgentKindNotFound, "agent kind not found")
	}
	return rowToAgentKind(row), nil
}

func (r *pgRepo) DeleteAgentKind(ctx context.Context, id uuid.UUID) error {
	if err := r.q.DeleteAgentKind(ctx, toPgUUID(id)); err != nil {
		if IsForeignKeyViolation(err) {
			return errs.New(errs.CodeAcpAgentKindInUse, "agent kind has active sessions")
		}
		return errs.Wrap(err, errs.CodeInternal, "delete agent kind")
	}
	return nil
}

func (r *pgRepo) CountAgentKindUsage(ctx context.Context, id uuid.UUID) (int64, error) {
	n, err := r.q.CountAgentKindUsage(ctx, toPgUUID(id))
	if err != nil {
		return 0, errs.Wrap(err, errs.CodeInternal, "count agent kind usage")
	}
	return n, nil
}

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

func rowToAgentKind(row acpsqlc.AcpAgentKind) *AgentKind {
	return &AgentKind{
		ID:           fromPgUUID(row.ID),
		Name:         row.Name,
		DisplayName:  row.DisplayName,
		Description:  row.Description,
		BinaryPath:   row.BinaryPath,
		Args:         row.Args,
		EncryptedEnv: row.EncryptedEnv,
		Enabled:      row.Enabled,
		CreatedBy:    fromPgUUID(row.CreatedBy),
		CreatedAt:    row.CreatedAt.Time,
		UpdatedAt:    row.UpdatedAt.Time,
	}
}

// ===== Session =====

func (r *pgRepo) InsertSession(ctx context.Context, s *Session) (*Session, error) {
	row, err := r.q.InsertSession(ctx, acpsqlc.InsertSessionParams{
		ID:             toPgUUID(s.ID),
		WorkspaceID:    toPgUUID(s.WorkspaceID),
		ProjectID:      toPgUUID(s.ProjectID),
		AgentKindID:    toPgUUID(s.AgentKindID),
		UserID:         toPgUUID(s.UserID),
		IssueID:        toPgUUIDPtr(s.IssueID),
		RequirementID:  toPgUUIDPtr(s.RequirementID),
		Branch:         s.Branch,
		CwdPath:        s.CwdPath,
		IsMainWorktree: s.IsMainWorktree,
		Status:         string(s.Status),
	})
	if err != nil {
		return nil, errs.Wrap(err, errs.CodeInternal, "insert session")
	}
	return rowToSession(row), nil
}

func (r *pgRepo) GetSessionByID(ctx context.Context, id uuid.UUID) (*Session, error) {
	row, err := r.q.GetSessionByID(ctx, toPgUUID(id))
	if err != nil {
		return nil, pgxNoRowsToErrNotFound(err, errs.CodeAcpSessionNotFound, "session not found")
	}
	return rowToSession(row), nil
}

func (r *pgRepo) ListSessionsByUser(ctx context.Context, userID uuid.UUID) ([]*Session, error) {
	rows, err := r.q.ListSessionsByUser(ctx, toPgUUID(userID))
	if err != nil {
		return nil, errs.Wrap(err, errs.CodeInternal, "list sessions by user")
	}
	out := make([]*Session, 0, len(rows))
	for _, row := range rows {
		out = append(out, rowToSession(row))
	}
	return out, nil
}

func (r *pgRepo) ListAllSessions(ctx context.Context) ([]*Session, error) {
	rows, err := r.q.ListAllSessions(ctx)
	if err != nil {
		return nil, errs.Wrap(err, errs.CodeInternal, "list all sessions")
	}
	out := make([]*Session, 0, len(rows))
	for _, row := range rows {
		out = append(out, rowToSession(row))
	}
	return out, nil
}

func (r *pgRepo) UpdateSessionRunning(ctx context.Context, id uuid.UUID, agentSessionID string, pid int32) error {
	if err := r.q.UpdateSessionRunning(ctx, acpsqlc.UpdateSessionRunningParams{
		ID:             toPgUUID(id),
		AgentSessionID: &agentSessionID,
		Pid:            &pid,
	}); err != nil {
		return errs.Wrap(err, errs.CodeInternal, "update session running")
	}
	return nil
}

func (r *pgRepo) UpdateSessionFinished(ctx context.Context, id uuid.UUID, status SessionStatus, exitCode *int32, lastErr *string) error {
	if err := r.q.UpdateSessionFinished(ctx, acpsqlc.UpdateSessionFinishedParams{
		ID:        toPgUUID(id),
		Status:    string(status),
		ExitCode:  exitCode,
		LastError: lastErr,
	}); err != nil {
		return errs.Wrap(err, errs.CodeInternal, "update session finished")
	}
	return nil
}

func (r *pgRepo) CountActiveSessions(ctx context.Context) (int64, error) {
	n, err := r.q.CountActiveSessions(ctx)
	if err != nil {
		return 0, errs.Wrap(err, errs.CodeInternal, "count active sessions")
	}
	return n, nil
}

func (r *pgRepo) CountActiveSessionsByUser(ctx context.Context, userID uuid.UUID) (int64, error) {
	n, err := r.q.CountActiveSessionsByUser(ctx, toPgUUID(userID))
	if err != nil {
		return 0, errs.Wrap(err, errs.CodeInternal, "count active sessions by user")
	}
	return n, nil
}

func (r *pgRepo) ResetStuckSessionsOnRestart(ctx context.Context) ([]*Session, error) {
	rows, err := r.q.ResetStuckSessionsOnRestart(ctx)
	if err != nil {
		return nil, errs.Wrap(err, errs.CodeInternal, "reset stuck sessions on restart")
	}
	out := make([]*Session, 0, len(rows))
	for _, row := range rows {
		out = append(out, &Session{
			ID:             fromPgUUID(row.ID),
			WorkspaceID:    fromPgUUID(row.WorkspaceID),
			UserID:         fromPgUUID(row.UserID),
			Branch:         row.Branch,
			AgentKindID:    fromPgUUID(row.AgentKindID),
			IsMainWorktree: row.IsMainWorktree,
			Status:         SessionCrashed,
		})
	}
	return out, nil
}

func (r *pgRepo) AcquireMainWorktree(ctx context.Context, sessionID, workspaceID uuid.UUID) (bool, error) {
	n, err := r.q.AcquireMainWorktree(ctx, acpsqlc.AcquireMainWorktreeParams{
		ActiveMainSessionID: toPgUUIDPtr(&sessionID),
		ID:                  toPgUUID(workspaceID),
	})
	if err != nil {
		return false, errs.Wrap(err, errs.CodeInternal, "acquire main worktree")
	}
	return n > 0, nil
}

func (r *pgRepo) ReleaseMainWorktree(ctx context.Context, workspaceID, sessionID uuid.UUID) (bool, error) {
	n, err := r.q.ReleaseMainWorktree(ctx, acpsqlc.ReleaseMainWorktreeParams{
		ID:                  toPgUUID(workspaceID),
		ActiveMainSessionID: toPgUUIDPtr(&sessionID),
	})
	if err != nil {
		return false, errs.Wrap(err, errs.CodeInternal, "release main worktree")
	}
	return n > 0, nil
}

func rowToSession(row acpsqlc.AcpSession) *Session {
	var endedAt *time.Time
	if row.EndedAt.Valid {
		t := row.EndedAt.Time
		endedAt = &t
	}
	return &Session{
		ID:             fromPgUUID(row.ID),
		WorkspaceID:    fromPgUUID(row.WorkspaceID),
		ProjectID:      fromPgUUID(row.ProjectID),
		AgentKindID:    fromPgUUID(row.AgentKindID),
		UserID:         fromPgUUID(row.UserID),
		IssueID:        fromPgUUIDPtr(row.IssueID),
		RequirementID:  fromPgUUIDPtr(row.RequirementID),
		AgentSessionID: row.AgentSessionID,
		Branch:         row.Branch,
		CwdPath:        row.CwdPath,
		IsMainWorktree: row.IsMainWorktree,
		Status:         SessionStatus(row.Status),
		PID:            row.Pid,
		ExitCode:       row.ExitCode,
		LastError:      row.LastError,
		StartedAt:      row.StartedAt.Time,
		EndedAt:        endedAt,
	}
}

// ===== Event =====

func (r *pgRepo) InsertEvent(ctx context.Context, e *Event) (*Event, error) {
	row, err := r.q.InsertEvent(ctx, acpsqlc.InsertEventParams{
		SessionID:   toPgUUID(e.SessionID),
		Direction:   string(e.Direction),
		RpcKind:     string(e.RPCKind),
		Method:      e.Method,
		Payload:     e.Payload,
		PayloadSize: e.PayloadSize,
		Truncated:   e.Truncated,
	})
	if err != nil {
		return nil, errs.Wrap(err, errs.CodeInternal, "insert acp event")
	}
	return rowToEvent(row), nil
}

func (r *pgRepo) ListEventsSince(ctx context.Context, sessionID uuid.UUID, sinceID int64, limit int32) ([]*Event, error) {
	rows, err := r.q.ListEventsSince(ctx, acpsqlc.ListEventsSinceParams{
		SessionID: toPgUUID(sessionID),
		ID:        sinceID,
		Limit:     limit,
	})
	if err != nil {
		return nil, errs.Wrap(err, errs.CodeInternal, "list events since")
	}
	out := make([]*Event, 0, len(rows))
	for _, row := range rows {
		out = append(out, rowToEvent(row))
	}
	return out, nil
}

func (r *pgRepo) PurgeEventsBefore(ctx context.Context, before time.Time) (int64, error) {
	n, err := r.q.PurgeEventsBefore(ctx, pgtype.Timestamptz{Time: before, Valid: true})
	if err != nil {
		return 0, errs.Wrap(err, errs.CodeInternal, "purge acp events")
	}
	return n, nil
}

func rowToEvent(row acpsqlc.AcpEvent) *Event {
	return &Event{
		ID:          row.ID,
		SessionID:   fromPgUUID(row.SessionID),
		Direction:   EventDirection(row.Direction),
		RPCKind:     RPCKind(row.RpcKind),
		Method:      row.Method,
		Payload:     row.Payload,
		PayloadSize: row.PayloadSize,
		Truncated:   row.Truncated,
		CreatedAt:   row.CreatedAt.Time,
	}
}