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

package codeindex

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"os"
	"path/filepath"

	"github.com/yan1h/agent-coding-workflow/internal/infra/git"
	"github.com/yan1h/agent-coding-workflow/internal/infra/llm"
	"github.com/yan1h/agent-coding-workflow/internal/jobs"
)

// EmbedBatchSize bounds how many chunks are sent to the embedder per request.
const EmbedBatchSize = 64

// FailureNotifier is an optional best-effort callback invoked when a build run
// fails, so app.go can broadcast to admins without the runner importing user/notify.
type FailureNotifier func(ctx context.Context, run *Run, cause error)

// BuildRunner is the jobs.Handler that materializes a code_index_runs row into
// embedded code_chunks. It is event-driven (enqueued by EnqueueBuild), never a
// periodic ticker. Idempotent: re-running for the same run replaces its chunks.
type BuildRunner struct {
	repo     Repository
	locator  WorkspaceLocator
	gitr     git.Runner
	embedSel EmbedderSelector
	chunker  *Chunker
	notify   FailureNotifier
	log      *slog.Logger
}

// NewBuildRunner constructs a BuildRunner. notify and log may be nil.
func NewBuildRunner(repo Repository, locator WorkspaceLocator, gitr git.Runner, embedSel EmbedderSelector, chunker *Chunker, notify FailureNotifier, log *slog.Logger) *BuildRunner {
	if chunker == nil {
		chunker = NewChunker()
	}
	if log == nil {
		log = slog.Default()
	}
	return &BuildRunner{repo: repo, locator: locator, gitr: gitr, embedSel: embedSel, chunker: chunker, notify: notify, log: log}
}

var _ jobs.Handler = (*BuildRunner)(nil)

// Type returns the job type this handler processes.
func (h *BuildRunner) Type() jobs.JobType { return JobTypeBuild }

// Handle builds the index for the run referenced in the payload. On any error it
// marks the run failed and notifies admins; the returned error drives jobs retry.
func (h *BuildRunner) Handle(ctx context.Context, job jobs.Job) error {
	var p buildPayload
	if err := json.Unmarshal(job.Payload, &p); err != nil {
		return jobs.Permanent(fmt.Errorf("codeindex.build: bad payload: %w", err))
	}
	run, err := h.repo.GetRun(ctx, p.RunID)
	if err != nil {
		return jobs.Permanent(fmt.Errorf("codeindex.build: load run %s: %w", p.RunID, err))
	}
	// Skip runs superseded by a newer one (stale) — no-op success.
	if run.Status == StatusStale || run.Status == StatusCompleted {
		return nil
	}

	if err := h.build(ctx, run); err != nil {
		if markErr := h.repo.MarkFailed(ctx, run.ID, truncErr(err)); markErr != nil {
			h.log.Error("codeindex.build.mark_failed", "run_id", run.ID, "err", markErr.Error())
		}
		if h.notify != nil {
			h.notify(ctx, run, err)
		}
		return err // retryable unless wrapped Permanent inside build
	}
	return nil
}

func (h *BuildRunner) build(ctx context.Context, run *Run) error {
	emb, _, model, dims, err := h.embedSel.Embedder(ctx)
	if err != nil {
		return jobs.Permanent(fmt.Errorf("codeindex.build: embedder unavailable: %w", err))
	}
	// Validate dims match the fixed pgvector column to avoid corrupt inserts.
	if d := emb.EmbedDims(model); d != 0 && d != llm.PlatformEmbeddingDims {
		return jobs.Permanent(fmt.Errorf("codeindex.build: embedder dims %d != platform %d", d, llm.PlatformEmbeddingDims))
	}
	if dims != llm.PlatformEmbeddingDims {
		return jobs.Permanent(fmt.Errorf("codeindex.build: run dims %d != platform %d", dims, llm.PlatformEmbeddingDims))
	}

	// Resolve the worktree/main directory holding the commit's files.
	dir := ""
	if run.WorktreeID != nil {
		path, _, _, lerr := h.locator.LocateWorktree(ctx, *run.WorktreeID)
		if lerr != nil {
			return fmt.Errorf("locate worktree: %w", lerr)
		}
		dir = path
	} else {
		info, lerr := h.locator.LocateWorkspace(ctx, run.WorkspaceID)
		if lerr != nil {
			return fmt.Errorf("locate workspace: %w", lerr)
		}
		dir = info.MainPath
	}

	if err := h.repo.MarkRunning(ctx, run.ID); err != nil {
		return fmt.Errorf("mark running: %w", err)
	}

	files, err := h.gitr.ListTrackedFiles(ctx, dir)
	if err != nil {
		return fmt.Errorf("git ls-files: %w", err)
	}

	var rows []ChunkRow
	filesIndexed := 0
	for _, rel := range files {
		abs := filepath.Join(dir, filepath.FromSlash(rel))
		fi, serr := os.Stat(abs)
		if serr != nil || fi.IsDir() {
			continue
		}
		if int(fi.Size()) > h.chunker.maxFileBytes() || isVendored(rel) {
			continue
		}
		content, rerr := os.ReadFile(abs)
		if rerr != nil {
			continue
		}
		chunks := h.chunker.Chunk(rel, content)
		if len(chunks) == 0 {
			continue
		}
		filesIndexed++
		for _, c := range chunks {
			rows = append(rows, ChunkRow{
				FilePath:   rel,
				Lang:       c.Lang,
				StartLine:  c.StartLine,
				EndLine:    c.EndLine,
				Content:    c.Content,
				ContentSHA: c.ContentSHA,
				TokenCount: llm.EstimateTokens(c.Content),
			})
		}
	}

	// Batch-embed chunk contents (input order preserved per batch).
	for start := 0; start < len(rows); start += EmbedBatchSize {
		end := start + EmbedBatchSize
		if end > len(rows) {
			end = len(rows)
		}
		inputs := make([]string, 0, end-start)
		for i := start; i < end; i++ {
			inputs = append(inputs, rows[i].Content)
		}
		vecs, eerr := emb.Embed(ctx, model, inputs)
		if eerr != nil {
			return fmt.Errorf("embed batch [%d:%d]: %w", start, end, eerr)
		}
		if len(vecs) != end-start {
			return fmt.Errorf("embed batch [%d:%d]: got %d vectors", start, end, len(vecs))
		}
		for i := start; i < end; i++ {
			rows[i].Embedding = vecs[i-start]
		}
	}

	if err := h.repo.ReplaceChunks(ctx, run.ID, run.WorkspaceID, run.CommitSHA, rows); err != nil {
		return fmt.Errorf("insert chunks: %w", err)
	}
	if err := h.repo.MarkCompleted(ctx, run.ID, filesIndexed, len(rows)); err != nil {
		return fmt.Errorf("mark completed: %w", err)
	}
	h.log.Info("codeindex.build.completed",
		"run_id", run.ID, "workspace_id", run.WorkspaceID, "commit", run.CommitSHA,
		"files", filesIndexed, "chunks", len(rows))
	return nil
}

func truncErr(err error) string {
	s := err.Error()
	if len(s) > 1000 {
		return s[:1000]
	}
	return s
}