Agent-System/internal/controlcenter/docs/docs.go

// Package docs syncs .doc markdown files into cc_documents/cc_doc_chunks and
// embeds chunk content for pgvector retrieval. Embedding failures degrade
// gracefully: rows are stored with NULL embeddings and the sync succeeds.
package docs

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io/fs"
	"log/slog"
	"os"
	"path/filepath"
	"strings"
	"time"

	"gopkg.in/yaml.v3"

	"chat-agent-system/internal/controlcenter/embed"
	"chat-agent-system/internal/controlcenter/store"
)

// Syncer ingests repository docs and stage artifacts into the vector store.
type Syncer struct {
	st       *store.Store
	repoRoot string
	emb      embed.Embedder // nil = embeddings unavailable, chunks stored with NULL
}

func New(st *store.Store, repoRoot string, emb embed.Embedder) *Syncer {
	return &Syncer{st: st, repoRoot: repoRoot, emb: emb}
}

// Stats summarizes one Sync run.
type Stats struct {
	Added       int  `json:"added"`
	Updated     int  `json:"updated"`
	Removed     int  `json:"removed"`
	Unchanged   int  `json:"unchanged"`
	Chunks      int  `json:"chunks"`
	EmbedFailed bool `json:"embed_failed"`
}

// Sync walks <repoRoot>/.doc for *.md files (skipping assets/ directories),
// upserts changed documents, re-chunks them idempotently, and removes rows
// for files that disappeared.
func (s *Syncer) Sync(ctx context.Context) (Stats, error) {
	var stats Stats
	docRoot := filepath.Join(s.repoRoot, ".doc")

	type diskDoc struct {
		rel  string // repo-relative slash path, e.g. .doc/ops/CURRENT.md
		sha  string
		full string
	}
	seen := map[string]diskDoc{}
	err := filepath.WalkDir(docRoot, func(path string, d fs.DirEntry, err error) error {
		if err != nil {
			return err
		}
		if d.IsDir() {
			if d.Name() == "assets" && path != docRoot {
				return filepath.SkipDir
			}
			return nil
		}
		if !strings.HasSuffix(d.Name(), ".md") {
			return nil
		}
		rel, err := filepath.Rel(s.repoRoot, path)
		if err != nil {
			return err
		}
		rel = filepath.ToSlash(rel)
		raw, err := os.ReadFile(path)
		if err != nil {
			return fmt.Errorf("read %s: %w", rel, err)
		}
		seen[rel] = diskDoc{rel: rel, sha: sha256hex(raw), full: path}
		return nil
	})
	if err != nil {
		if os.IsNotExist(err) {
			slog.Warn("doc sync: .doc root missing", "root", docRoot)
			return stats, nil
		}
		return stats, err
	}

	// Existing rows, to diff against.
	existing := map[string]struct {
		id  string
		sha string
	}{}
	rows, err := s.st.Pool.Query(ctx, `SELECT id, path, sha256 FROM cc_documents`)
	if err != nil {
		return stats, fmt.Errorf("list cc_documents: %w", err)
	}
	for rows.Next() {
		var id, path, sha string
		if err := rows.Scan(&id, &path, &sha); err != nil {
			rows.Close()
			return stats, err
		}
		existing[path] = struct{ id, sha string }{id, sha}
	}
	rows.Close()

	// Remove documents whose files disappeared (chunks cascade).
	for path, ex := range existing {
		if _, ok := seen[path]; ok {
			continue
		}
		if _, err := s.st.Pool.Exec(ctx, `DELETE FROM cc_documents WHERE id = $1`, ex.id); err != nil {
			return stats, fmt.Errorf("delete vanished doc %s: %w", path, err)
		}
		stats.Removed++
	}

	for _, d := range seen {
		if ex, ok := existing[d.rel]; ok && ex.sha == d.sha {
			stats.Unchanged++
			continue
		}
		raw, err := os.ReadFile(d.full)
		if err != nil {
			return stats, fmt.Errorf("read %s: %w", d.rel, err)
		}
		failed, n, err := s.syncOne(ctx, d.rel, d.sha, raw)
		if err != nil {
			return stats, err
		}
		stats.EmbedFailed = stats.EmbedFailed || failed
		stats.Chunks += n
		if _, ok := existing[d.rel]; ok {
			stats.Updated++
		} else {
			stats.Added++
		}
	}
	return stats, nil
}

// syncOne upserts one document row and re-chunks it. Returns
// (embedFailed, newChunkCount, err).
func (s *Syncer) syncOne(ctx context.Context, rel, sha string, raw []byte) (bool, int, error) {
	fm, body := splitFrontMatter(string(raw))
	title := docTitle(fm, body, rel)
	fmJSON, _ := json.Marshal(fm)

	var docID string
	err := s.st.Pool.QueryRow(ctx, `
		INSERT INTO cc_documents (path, title, taxonomy, frontmatter, sha256, synced_at)
		VALUES ($1,$2,$3,$4,$5,now())
		ON CONFLICT (path) DO UPDATE SET
			title = EXCLUDED.title, taxonomy = EXCLUDED.taxonomy,
			frontmatter = EXCLUDED.frontmatter, sha256 = EXCLUDED.sha256,
			synced_at = now()
		RETURNING id`,
		rel, title, taxonomyOf(rel), fmJSON, sha).Scan(&docID)
	if err != nil {
		return false, 0, fmt.Errorf("upsert doc %s: %w", rel, err)
	}
	failed, n, err := s.rechunk(ctx, rel, body, "doc", docID, "")
	return failed, n, err
}

// IndexArtifact chunks and embeds one cc_stage_artifacts row
// (doc_id NULL, chunk_hash prefixed "artifact:").
func (s *Syncer) IndexArtifact(ctx context.Context, artifactID string) error {
	var content string
	err := s.st.Pool.QueryRow(ctx,
		`SELECT content_md FROM cc_stage_artifacts WHERE id = $1`, artifactID).Scan(&content)
	if err != nil {
		return fmt.Errorf("load artifact %s: %w", artifactID, err)
	}
	failed, _, err := s.rechunk(ctx, artifactID, content, "artifact", "", artifactID)
	if failed {
		slog.Warn("artifact indexed without embeddings", "artifact", artifactID)
	}
	return err
}

// rechunk diffs chunk_hash sets for an owner (doc or artifact) and inserts
// only new chunks; unchanged chunks keep their embeddings.
// ownerKind is "doc" or "artifact"; exactly one of docID/artifactID is set.
func (s *Syncer) rechunk(ctx context.Context, sourceRef, content, ownerKind, docID, artifactID string) (bool, int, error) {
	chunks := chunkMarkdown(content)

	existing := map[string]bool{}
	q := `SELECT chunk_hash FROM cc_doc_chunks WHERE doc_id = $1`
	arg := any(docID)
	if ownerKind == "artifact" {
		q = `SELECT chunk_hash FROM cc_doc_chunks WHERE artifact_id = $1`
		arg = artifactID
	}
	rows, err := s.st.Pool.Query(ctx, q, arg)
	if err != nil {
		return false, 0, fmt.Errorf("list chunks: %w", err)
	}
	for rows.Next() {
		var h string
		if err := rows.Scan(&h); err != nil {
			rows.Close()
			return false, 0, err
		}
		existing[h] = true
	}
	rows.Close()

	want := make(map[string]Chunk, len(chunks))
	var fresh []Chunk
	for _, c := range chunks {
		h := chunkHash(ownerKind, sourceRef, c.Seq, c.Content)
		want[h] = c
		if !existing[h] {
			fresh = append(fresh, c)
		}
	}

	// Drop chunks that are no longer part of the source.
	for h := range existing {
		if _, ok := want[h]; !ok {
			if _, err := s.st.Pool.Exec(ctx, `DELETE FROM cc_doc_chunks WHERE chunk_hash = $1`, h); err != nil {
				return false, 0, fmt.Errorf("delete stale chunk: %w", err)
			}
		}
	}
	if len(fresh) == 0 {
		return false, 0, nil
	}

	texts := make([]string, len(fresh))
	for i, c := range fresh {
		texts[i] = c.Content
	}
	var vecs [][]float32
	embedFailed := false
	if s.emb != nil {
		vecs, err = s.emb.Embed(ctx, texts)
		if err != nil {
			slog.Warn("embedding unavailable; chunks stored without vectors",
				"owner", ownerKind, "ref", sourceRef, "chunks", len(fresh), "err", err)
			embedFailed = true
			vecs = nil
		}
	} else {
		embedFailed = true
	}

	inserted := 0
	for i, c := range fresh {
		h := chunkHash(ownerKind, sourceRef, c.Seq, c.Content)
		var vecLit any
		if !embedFailed && i < len(vecs) {
			vecLit = embed.VectorLiteral(vecs[i])
		}
		var dID, aID any
		if ownerKind == "doc" {
			dID = docID
		} else {
			aID = artifactID
		}
		_, err := s.st.Pool.Exec(ctx, `
			INSERT INTO cc_doc_chunks (doc_id, artifact_id, seq, heading, content, chunk_hash, embedding)
			VALUES ($1,$2,$3,$4,$5,$6, $7::vector)
			ON CONFLICT (chunk_hash) DO NOTHING`,
			dID, aID, c.Seq, c.Heading, c.Content, h, vecLit)
		if err != nil {
			return embedFailed, inserted, fmt.Errorf("insert chunk %d: %w", c.Seq, err)
		}
		inserted++
	}
	return embedFailed, inserted, nil
}

// chunkHash is stable per (owner, source, seq, content); "artifact:" prefix
// keeps artifact chunks distinct from document chunks.
func chunkHash(ownerKind, sourceRef string, seq int, content string) string {
	prefix := ""
	if ownerKind == "artifact" {
		prefix = "artifact:"
	}
	return sha256hex([]byte(fmt.Sprintf("%s%s:%d:%s", prefix, sourceRef, seq, content)))
}

// StartPeriodic runs Sync immediately, then every interval until ctx is
// cancelled (same ticker pattern as the lease reaper).
func (s *Syncer) StartPeriodic(ctx context.Context, interval time.Duration) {
	go func() {
		s.runOnce(ctx)
		t := time.NewTicker(interval)
		defer t.Stop()
		for {
			select {
			case <-ctx.Done():
				return
			case <-t.C:
				s.runOnce(ctx)
			}
		}
	}()
}

func (s *Syncer) runOnce(ctx context.Context) {
	stats, err := s.Sync(ctx)
	if err != nil {
		slog.Error("doc sync", "err", err)
		return
	}
	slog.Info("doc sync done",
		"added", stats.Added, "updated", stats.Updated, "removed", stats.Removed,
		"unchanged", stats.Unchanged, "chunks", stats.Chunks, "embed_failed", stats.EmbedFailed)
}

// splitFrontMatter extracts a leading --- yaml --- block; returns the parsed
// map (possibly empty) and the remaining body.
func splitFrontMatter(md string) (map[string]any, string) {
	fm := map[string]any{}
	body := md
	if rest, ok := strings.CutPrefix(md, "---\n"); ok {
		if end := strings.Index(rest, "\n---"); end >= 0 {
			raw := rest[:end]
			after := rest[end+len("\n---"):]
			if after == "" || after[0] == '\n' || after[0] == '\r' {
				if err := yaml.Unmarshal([]byte(raw), &fm); err != nil {
					slog.Warn("doc sync: bad front matter", "err", err)
					fm = map[string]any{}
				} else {
					body = strings.TrimLeft(after, "\r\n")
				}
			}
		}
	}
	if fm == nil {
		fm = map[string]any{}
	}
	return fm, body
}

// docTitle prefers front matter title, then the first level-1 heading, then
// the file name.
func docTitle(fm map[string]any, body, rel string) string {
	if t, ok := fm["title"].(string); ok && strings.TrimSpace(t) != "" {
		return strings.TrimSpace(t)
	}
	for _, line := range strings.Split(body, "\n") {
		if h, ok := strings.CutPrefix(strings.TrimSpace(line), "# "); ok {
			if t := strings.TrimSpace(h); t != "" {
				return t
			}
		}
	}
	base := filepath.Base(rel)
	return strings.TrimSuffix(base, filepath.Ext(base))
}

// taxonomyOf returns the first directory segment under .doc/ ("" for files
// directly at the root).
func taxonomyOf(rel string) string {
	rest, ok := strings.CutPrefix(rel, ".doc/")
	if !ok {
		return ""
	}
	if i := strings.IndexByte(rest, '/'); i > 0 {
		return rest[:i]
	}
	return ""
}

func sha256hex(b []byte) string {
	sum := sha256.Sum256(b)
	return hex.EncodeToString(sum[:])
}