diff --git a/internal/acp/repository.go b/internal/acp/repository.go new file mode 100644 index 0000000..ecb069b --- /dev/null +++ b/internal/acp/repository.go @@ -0,0 +1,436 @@ +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) +} + +// 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)} +} + +// ===== 类型转换 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, + } +}