kenn-io--agentsview/internal/vector/manager.go

package vector

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"time"

	kitvec "go.kenn.io/kit/vector"
	"go.kenn.io/kit/vector/sqlitevec"
)

// ErrBuildRunning is returned by StartBuild when a build is already in
// flight, and by Activate/Retire when they are called while one is running.
var ErrBuildRunning = errors.New("an embeddings build is already running")

// ErrGenerationRefused indicates Activate or Retire declined to change a
// generation's state without --force (incomplete coverage, or retiring the
// active generation). Match it with errors.Is; the error's own message
// carries the specific, user-facing reason rather than a fixed string.
var ErrGenerationRefused = errors.New("generation state change refused")

// refusalError carries a specific, literal refusal message while still
// satisfying errors.Is(err, ErrGenerationRefused), so callers get an
// unprefixed message (e.g. for direct display) alongside a stable sentinel
// for status-code mapping.
type refusalError struct {
	msg string
}

func refusedf(format string, args ...any) error {
	return &refusalError{msg: fmt.Sprintf(format, args...)}
}

func (e *refusalError) Error() string { return e.msg }

func (e *refusalError) Is(target error) bool { return target == ErrGenerationRefused }

// ErrUnknownServer indicates a build request named an embeddings server that
// is not defined in the manager's encoder set — caller input, not a manager
// fault. Match it with errors.Is; the error's own message carries the
// offending name.
var ErrUnknownServer = errors.New("unknown embeddings server")

type unknownServerError struct {
	name string
}

func (e *unknownServerError) Error() string {
	return fmt.Sprintf("no embeddings server named %q", e.name)
}

func (e *unknownServerError) Is(target error) bool { return target == ErrUnknownServer }

// Manager serializes embedding builds over one Index: only one Build call
// may run at a time, whether triggered via StartBuild (async, for the HTTP
// API) or TryBuild (sync, for a periodic scheduler). Activate and Retire
// are likewise serialized against each other and against build starts, so
// their check-then-act refusal invariants (Missing coverage, the
// active-generation check) hold under concurrent calls.
type Manager struct {
	ix       *Index
	src      UnitSource
	encoders EncoderSet
	gen      kitvec.Generation

	// opMu serializes lifecycle operations: build starts (begin) and the
	// whole of Activate/Retire. It is never held across a running build —
	// begin releases it once running is set — so StartBuild stays
	// non-blocking while a build is in flight.
	opMu sync.Mutex

	// mu guards running and status; held only for short field updates so
	// Status() stays responsive during a build.
	mu      sync.Mutex
	running bool
	status  BuildStatus

	// now stamps BuildStatus.StartedAt when a build begins; a test hook
	// defaulting to time.Now.
	now func() time.Time
}

// BuildRequest is the caller-controlled subset of BuildOptions the manager
// exposes; encode settings and Progress are the manager's own concerns.
type BuildRequest struct {
	FullRebuild bool `json:"full_rebuild,omitempty"`
	Backstop    bool `json:"backstop,omitempty"`
	// IncludeAutomated is the resolved include-automated scope for this
	// build (caller-resolved from config and, for the CLI's one-off
	// --include-automated flag, its override). See BuildOptions.IncludeAutomated.
	IncludeAutomated bool `json:"include_automated,omitempty"`
	// Using names the embeddings server (an EncoderSet entry) this build
	// encodes against; empty selects the set's default. Server choice is
	// transport only — every server encodes the same model, so it never
	// affects the generation fingerprint.
	Using string `json:"using,omitempty"`
}

// BuildStatus reports the manager's current build state, for polling
// clients (CLI and HTTP API).
type BuildStatus struct {
	Running bool `json:"running"`
	// BuildID identifies one build within this daemon process: it increments
	// each time a build starts, so a polling client can tell two builds
	// apart (e.g. to reset a progress-rate estimator across page reloads)
	// instead of treating unrelated builds as continuous. Zero until the
	// first build of the process.
	BuildID int64 `json:"build_id,omitempty"`
	// StartedAt is when the current (or, once it finished, most recent)
	// build began, RFC3339 UTC. Empty until the first build of the process.
	StartedAt  string       `json:"started_at,omitempty"`
	Phase      string       `json:"phase,omitempty"`
	Done       int64        `json:"done"`
	Total      int64        `json:"total"`
	LastError  string       `json:"last_error,omitempty"`
	LastResult *BuildResult `json:"last_result,omitempty"`
	// Model and Dimension identify the configured embedding space the
	// manager builds ([vector.embeddings] model/dimension), so status
	// consumers can display the target space even before any generation
	// exists. They describe config, not a specific build.
	Model     string `json:"model,omitempty"`
	Dimension int    `json:"dimension,omitempty"`
}

// EncodeSettings groups the encode-shape knobs of one embeddings server,
// resolved from config ([vector.embeddings.servers.<name>] batch_size and
// concurrency).
type EncodeSettings struct {
	// BatchSize is the number of inputs sent per HTTP call.
	BatchSize int
	// Concurrency is the number of documents encoded in parallel.
	Concurrency int
}

// ManagedEncoder pairs one embeddings server's encoder with the encode
// settings tuned for that server.
type ManagedEncoder struct {
	Encode   kitvec.EncodeFunc
	Settings EncodeSettings
}

// EncoderSet is the named embeddings servers a Manager can build with. All
// entries encode the same model — the embedding-space identity is global
// config — so a build may use any of them interchangeably; Default names
// the entry used when a BuildRequest doesn't select one.
type EncoderSet struct {
	Default string
	ByName  map[string]ManagedEncoder
}

// NewManager creates a Manager that builds gen's embedding space over ix,
// scanning src and encoding with one of encoders' entries per build (the
// default, or BuildRequest.Using). Each encoder is wrapped so a panic
// inside it surfaces as an encode error rather than crashing the process
// (see recoveringEncoder).
func NewManager(
	ix *Index, src UnitSource, encoders EncoderSet, gen kitvec.Generation,
) *Manager {
	wrapped := EncoderSet{Default: encoders.Default, ByName: make(map[string]ManagedEncoder, len(encoders.ByName))}
	for name, me := range encoders.ByName {
		me.Encode = recoveringEncoder(me.Encode)
		wrapped.ByName[name] = me
	}
	return &Manager{ix: ix, src: src, encoders: wrapped, gen: gen, now: time.Now}
}

// resolveEncoder picks the encoder a build request encodes with: the named
// entry when Using is set, the set's default otherwise. It fails when the
// name is unknown so a mistyped --using errors before a build starts.
func (m *Manager) resolveEncoder(req BuildRequest) (ManagedEncoder, error) {
	name := req.Using
	if name == "" {
		name = m.encoders.Default
	}
	me, ok := m.encoders.ByName[name]
	if !ok {
		return ManagedEncoder{}, &unknownServerError{name: name}
	}
	return me, nil
}

// recoveringEncoder converts a panic in enc (a caller-supplied network
// client) into an ordinary encode error. This must wrap the encoder itself
// rather than rely on runBuild's recover: kit's EncodeBatched invokes
// encoders on its own worker goroutines, where a recover on the manager's
// build goroutine cannot reach.
func recoveringEncoder(enc kitvec.EncodeFunc) kitvec.EncodeFunc {
	return func(ctx context.Context, texts []string) (vectors [][]float32, err error) {
		defer func() {
			if r := recover(); r != nil {
				vectors = nil
				err = fmt.Errorf("encoder panicked: %v", r)
			}
		}()
		return enc(ctx, texts)
	}
}

// StartBuild launches a Build in a background goroutine, returning
// ErrBuildRunning if one is already in flight rather than queuing behind it.
// The goroutine runs against context.Background() so it outlives the HTTP
// request that triggered it.
func (m *Manager) StartBuild(req BuildRequest) error {
	me, err := m.resolveEncoder(req)
	if err != nil {
		return err
	}
	if err := m.begin(); err != nil {
		return err
	}
	go func() {
		result, err := m.runBuild(context.Background(), req, me)
		m.finish(result, err)
	}()
	return nil
}

// TryBuild runs one Build synchronously, for a periodic scheduler that
// should drop a scheduled run rather than queue it: it returns (false, nil)
// without starting anything if a build is already running.
func (m *Manager) TryBuild(ctx context.Context, req BuildRequest) (bool, error) {
	me, err := m.resolveEncoder(req)
	if err != nil {
		return false, err
	}
	if err := m.begin(); err != nil {
		return false, nil
	}
	result, err := m.runBuild(ctx, req, me)
	m.finish(result, err)
	return true, err
}

// Status returns a snapshot of the manager's current build state, stamped
// with the configured embedding space's identity.
func (m *Manager) Status() BuildStatus {
	m.mu.Lock()
	defer m.mu.Unlock()
	status := m.status
	if status.LastResult != nil {
		result := *status.LastResult
		status.LastResult = &result
	}
	status.Model = m.gen.Model
	status.Dimension = m.gen.Dimensions
	return status
}

// Generations delegates to the underlying Index, listing every generation
// with its coverage of the current mirror.
func (m *Manager) Generations(ctx context.Context) ([]GenerationInfo, error) {
	return m.ix.Generations(ctx)
}

// Activate transitions the generation identified by id to active, retiring
// whichever generation was previously active (in one transaction, via the
// same activateGeneration primitive Build's auto-activation uses, so two
// generations can never end up active simultaneously). Without force, it
// refuses when id's generation has documents still needing embedding
// (Missing > 0) or while a build is running. Serialized against Retire,
// other Activate calls, and build starts via opMu.
func (m *Manager) Activate(ctx context.Context, id int64, force bool) error {
	m.opMu.Lock()
	defer m.opMu.Unlock()
	if m.isRunning() {
		return ErrBuildRunning
	}

	target, err := m.ix.GenerationByID(ctx, id)
	if err != nil {
		return err
	}
	if !force && target.Missing > 0 {
		return refusedf("generation %d still has %d documents needing embedding; use --force",
			id, target.Missing)
	}
	return m.ix.activateGeneration(ctx, target.Fingerprint)
}

// Retire transitions the generation identified by id to retired. Without
// force, it refuses when id is the active generation or while a build is
// running. Serialized against Activate, other Retire calls, and build
// starts via opMu.
func (m *Manager) Retire(ctx context.Context, id int64, force bool) error {
	m.opMu.Lock()
	defer m.opMu.Unlock()
	if m.isRunning() {
		return ErrBuildRunning
	}

	target, err := m.ix.GenerationByID(ctx, id)
	if err != nil {
		return err
	}
	if !force && target.State == string(sqlitevec.StateActive) {
		return refusedf("generation %d is active; use --force to retire it", id)
	}
	return m.ix.SetStateByID(ctx, id, sqlitevec.StateRetired)
}

// begin transitions the manager into the running state, resetting the
// progress fields of status for the new run. It returns ErrBuildRunning
// without changing anything if a build is already in flight. Taking opMu
// first serializes build starts behind any in-flight Activate/Retire.
func (m *Manager) begin() error {
	m.opMu.Lock()
	defer m.opMu.Unlock()
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.running {
		return ErrBuildRunning
	}
	m.running = true
	m.status.Running = true
	m.status.BuildID++
	m.status.StartedAt = m.now().UTC().Format(time.RFC3339)
	m.status.Phase = ""
	m.status.Done = 0
	m.status.Total = 0
	return nil
}

// runBuild performs the actual Index.Build call, wiring the manager's
// reportProgress method in as the BuildOptions.Progress callback so Status
// reflects incremental progress while the build is in flight. It converts
// a panic (e.g. from the caller-supplied encoder, a network client) into an
// error so StartBuild's detached goroutine can never crash the process and
// TryBuild's caller sees a failure rather than a propagating panic; either
// way finish records it in LastError and clears the running state.
func (m *Manager) runBuild(
	ctx context.Context, req BuildRequest, me ManagedEncoder,
) (result BuildResult, err error) {
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("build panicked: %v", r)
		}
	}()
	return m.ix.Build(ctx, m.src, me.Encode, m.gen, BuildOptions{
		FullRebuild:      req.FullRebuild,
		Backstop:         req.Backstop,
		IncludeAutomated: req.IncludeAutomated,
		BatchSize:        me.Settings.BatchSize,
		Concurrency:      me.Settings.Concurrency,
		Progress:         m.reportProgress,
	})
}

// reportProgress updates status's progress fields under the manager's lock;
// it is passed as BuildOptions.Progress and so runs on the build goroutine.
func (m *Manager) reportProgress(p BuildProgress) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.status.Phase = p.Phase
	m.status.Done = p.Done
	m.status.Total = p.Total
}

// finish records a completed build's outcome and clears the running state.
// A successful build sets LastResult and clears any previous LastError; a
// failed build sets LastError and leaves the last successful LastResult (if
// any) untouched.
func (m *Manager) finish(result BuildResult, err error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.running = false
	m.status.Running = false
	if err != nil {
		m.status.LastError = err.Error()
		return
	}
	m.status.LastError = ""
	r := result
	m.status.LastResult = &r
}

func (m *Manager) isRunning() bool {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.running
}