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

package codeindex

import (
	"context"
	"errors"
	"time"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgtype"
	"github.com/jackc/pgx/v5/pgxpool"
	"github.com/pgvector/pgvector-go"

	codeindexsqlc "github.com/yan1h/agent-coding-workflow/internal/codeindex/sqlc"
	"github.com/yan1h/agent-coding-workflow/internal/infra/errs"
)

// ChunkRow is a fully-populated chunk ready for bulk insert (embedding included).
type ChunkRow struct {
	FilePath   string
	Lang       string
	StartLine  int
	EndLine    int
	Content    string
	ContentSHA []byte
	TokenCount int
	Embedding  []float32
}

// Repository is the codeindex persistence contract.
type Repository interface {
	InsertRun(ctx context.Context, r *Run) (*Run, error)
	GetRun(ctx context.Context, id uuid.UUID) (*Run, error)
	GetRunByCommit(ctx context.Context, wsID uuid.UUID, commitSHA, model string) (*Run, error)
	GetLatestCompletedRun(ctx context.Context, wsID uuid.UUID) (*Run, error)
	MarkRunning(ctx context.Context, id uuid.UUID) error
	MarkCompleted(ctx context.Context, id uuid.UUID, files, chunks int) error
	MarkFailed(ctx context.Context, id uuid.UUID, errMsg string) error
	MarkStale(ctx context.Context, wsID uuid.UUID) error
	MarkStuckStale(ctx context.Context, runningBefore time.Time) ([]uuid.UUID, error)
	// ReplaceChunks deletes any existing chunks for run then bulk-inserts rows.
	ReplaceChunks(ctx context.Context, runID, wsID uuid.UUID, commitSHA string, rows []ChunkRow) error
	// SearchKNN returns the top-k chunks of the latest completed run for wsID
	// ordered by cosine similarity to queryVec, with optional path/lang filters.
	SearchKNN(ctx context.Context, wsID uuid.UUID, queryVec []float32, topK int, pathPrefix, lang string) ([]Snippet, error)
}

// PgRepository is the production Repository on a pgxpool.Pool.
type PgRepository struct {
	pool *pgxpool.Pool
	q    *codeindexsqlc.Queries
}

// NewPostgresRepository builds a PgRepository.
func NewPostgresRepository(pool *pgxpool.Pool) *PgRepository {
	return &PgRepository{pool: pool, q: codeindexsqlc.New(pool)}
}

var _ Repository = (*PgRepository)(nil)

func pgUUID(u uuid.UUID) pgtype.UUID { return pgtype.UUID{Bytes: u, Valid: true} }

func pgUUIDPtr(u *uuid.UUID) pgtype.UUID {
	if u == nil {
		return pgtype.UUID{}
	}
	return pgtype.UUID{Bytes: *u, Valid: true}
}

func uuidPtr(p pgtype.UUID) *uuid.UUID {
	if !p.Valid {
		return nil
	}
	id := uuid.UUID(p.Bytes)
	return &id
}

func runFromRow(r codeindexsqlc.CodeIndexRun) *Run {
	out := &Run{
		ID:                  uuid.UUID(r.ID.Bytes),
		WorkspaceID:         uuid.UUID(r.WorkspaceID.Bytes),
		WorktreeID:          uuidPtr(r.WorktreeID),
		Branch:              r.Branch,
		CommitSHA:           r.CommitSha,
		Status:              r.Status,
		EmbeddingEndpointID: uuidPtr(r.EmbeddingEndpointID),
		EmbeddingModel:      r.EmbeddingModel,
		Dims:                int(r.Dims),
		FilesIndexed:        int(r.FilesIndexed),
		ChunksIndexed:       int(r.ChunksIndexed),
		Error:               r.Error,
	}
	if r.CreatedAt.Valid {
		out.CreatedAt = r.CreatedAt.Time
	}
	return out
}

func (r *PgRepository) InsertRun(ctx context.Context, run *Run) (*Run, error) {
	row, err := r.q.InsertRun(ctx, codeindexsqlc.InsertRunParams{
		ID:                  pgUUID(run.ID),
		WorkspaceID:         pgUUID(run.WorkspaceID),
		WorktreeID:          pgUUIDPtr(run.WorktreeID),
		Branch:              run.Branch,
		CommitSha:           run.CommitSHA,
		Status:              StatusPending,
		EmbeddingEndpointID: pgUUIDPtr(run.EmbeddingEndpointID),
		EmbeddingModel:      run.EmbeddingModel,
		Dims:                int32(run.Dims),
	})
	if err != nil {
		return nil, errs.Wrap(err, errs.CodeInternal, "insert code index run")
	}
	return runFromRow(row), nil
}

func (r *PgRepository) GetRun(ctx context.Context, id uuid.UUID) (*Run, error) {
	row, err := r.q.GetRun(ctx, pgUUID(id))
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, errs.New(errs.CodeNotFound, "code index run not found")
		}
		return nil, errs.Wrap(err, errs.CodeInternal, "get code index run")
	}
	return runFromRow(row), nil
}

func (r *PgRepository) GetRunByCommit(ctx context.Context, wsID uuid.UUID, commitSHA, model string) (*Run, error) {
	row, err := r.q.GetRunByCommit(ctx, codeindexsqlc.GetRunByCommitParams{
		WorkspaceID:    pgUUID(wsID),
		CommitSha:      commitSHA,
		EmbeddingModel: model,
	})
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, nil
		}
		return nil, errs.Wrap(err, errs.CodeInternal, "get run by commit")
	}
	return runFromRow(row), nil
}

func (r *PgRepository) GetLatestCompletedRun(ctx context.Context, wsID uuid.UUID) (*Run, error) {
	row, err := r.q.GetLatestCompletedRun(ctx, pgUUID(wsID))
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, nil
		}
		return nil, errs.Wrap(err, errs.CodeInternal, "get latest completed run")
	}
	return runFromRow(row), nil
}

func (r *PgRepository) MarkRunning(ctx context.Context, id uuid.UUID) error {
	if err := r.q.MarkRunRunning(ctx, pgUUID(id)); err != nil {
		return errs.Wrap(err, errs.CodeInternal, "mark run running")
	}
	return nil
}

func (r *PgRepository) MarkCompleted(ctx context.Context, id uuid.UUID, files, chunks int) error {
	if err := r.q.MarkRunCompleted(ctx, codeindexsqlc.MarkRunCompletedParams{
		ID:            pgUUID(id),
		FilesIndexed:  int32(files),
		ChunksIndexed: int32(chunks),
	}); err != nil {
		return errs.Wrap(err, errs.CodeInternal, "mark run completed")
	}
	return nil
}

func (r *PgRepository) MarkFailed(ctx context.Context, id uuid.UUID, errMsg string) error {
	msg := errMsg
	if err := r.q.MarkRunFailed(ctx, codeindexsqlc.MarkRunFailedParams{
		ID:    pgUUID(id),
		Error: &msg,
	}); err != nil {
		return errs.Wrap(err, errs.CodeInternal, "mark run failed")
	}
	return nil
}

func (r *PgRepository) MarkStale(ctx context.Context, wsID uuid.UUID) error {
	if err := r.q.MarkRunsStale(ctx, pgUUID(wsID)); err != nil {
		return errs.Wrap(err, errs.CodeInternal, "mark runs stale")
	}
	return nil
}

func (r *PgRepository) MarkStuckStale(ctx context.Context, runningBefore time.Time) ([]uuid.UUID, error) {
	rows, err := r.q.MarkStuckRunsStale(ctx, pgtype.Timestamptz{Time: runningBefore, Valid: true})
	if err != nil {
		return nil, errs.Wrap(err, errs.CodeInternal, "mark stuck runs stale")
	}
	out := make([]uuid.UUID, 0, len(rows))
	for _, p := range rows {
		out = append(out, uuid.UUID(p.Bytes))
	}
	return out, nil
}

// ReplaceChunks deletes prior chunks for the run and bulk-inserts the given rows
// in a single transaction. The embedding column uses the pgvector text literal.
func (r *PgRepository) ReplaceChunks(ctx context.Context, runID, wsID uuid.UUID, commitSHA string, rows []ChunkRow) error {
	tx, err := r.pool.Begin(ctx)
	if err != nil {
		return errs.Wrap(err, errs.CodeInternal, "begin chunk tx")
	}
	defer func() { _ = tx.Rollback(ctx) }()

	if _, err := tx.Exec(ctx, `DELETE FROM code_chunks WHERE run_id = $1`, pgUUID(runID)); err != nil {
		return errs.Wrap(err, errs.CodeInternal, "delete prior chunks")
	}

	const insertSQL = `INSERT INTO code_chunks (
		id, run_id, workspace_id, commit_sha, file_path, lang,
		start_line, end_line, content, content_sha, token_count, embedding
	) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12::vector)`

	batch := &pgx.Batch{}
	for _, row := range rows {
		var lang *string
		if row.Lang != "" {
			l := row.Lang
			lang = &l
		}
		var tok *int32
		if row.TokenCount > 0 {
			t := int32(row.TokenCount)
			tok = &t
		}
		batch.Queue(insertSQL,
			pgUUID(uuid.New()), pgUUID(runID), pgUUID(wsID), commitSHA,
			row.FilePath, lang, int32(row.StartLine), int32(row.EndLine),
			row.Content, row.ContentSHA, tok, pgvector.NewVector(row.Embedding),
		)
	}
	br := tx.SendBatch(ctx, batch)
	for range rows {
		if _, err := br.Exec(); err != nil {
			_ = br.Close()
			return errs.Wrap(err, errs.CodeInternal, "insert chunk")
		}
	}
	if err := br.Close(); err != nil {
		return errs.Wrap(err, errs.CodeInternal, "close chunk batch")
	}
	if err := tx.Commit(ctx); err != nil {
		return errs.Wrap(err, errs.CodeInternal, "commit chunk tx")
	}
	return nil
}

// SearchKNN runs cosine-distance KNN scoped to the latest completed run for the
// workspace. Score is 1 - cosine_distance (so higher is more similar). Optional
// path_prefix / lang filters narrow the candidate set.
func (r *PgRepository) SearchKNN(ctx context.Context, wsID uuid.UUID, queryVec []float32, topK int, pathPrefix, lang string) ([]Snippet, error) {
	if topK <= 0 {
		topK = 10
	}
	q := `
WITH latest AS (
  SELECT id FROM code_index_runs
  WHERE workspace_id = $1 AND status = 'completed'
  ORDER BY created_at DESC LIMIT 1
)
SELECT c.file_path, c.start_line, c.end_line, c.content, c.commit_sha,
       1 - (c.embedding <=> $2::vector) AS score
FROM code_chunks c
JOIN latest ON c.run_id = latest.id
WHERE ($3::text = '' OR c.file_path LIKE $3 || '%')
  AND ($4::text = '' OR c.lang = $4)
ORDER BY c.embedding <=> $2::vector
LIMIT $5`
	rows, err := r.pool.Query(ctx, q, pgUUID(wsID), pgvector.NewVector(queryVec), pathPrefix, lang, int32(topK))
	if err != nil {
		return nil, errs.Wrap(err, errs.CodeInternal, "knn search")
	}
	defer rows.Close()
	var out []Snippet
	for rows.Next() {
		var s Snippet
		var start, end int32
		if err := rows.Scan(&s.FilePath, &start, &end, &s.Content, &s.CommitSHA, &s.Score); err != nil {
			return nil, errs.Wrap(err, errs.CodeInternal, "scan knn row")
		}
		s.StartLine, s.EndLine = int(start), int(end)
		out = append(out, s)
	}
	if err := rows.Err(); err != nil {
		return nil, errs.Wrap(err, errs.CodeInternal, "knn rows")
	}
	return out, nil
}