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 }