llmgw/internal/config/layer.go

package config

import (
	"bytes"
	"errors"
	"fmt"
	"maps"
	"os"
	"reflect"
	"slices"
	"strings"

	"go.yaml.in/yaml/v3"
)

// localKeys are the keys of server.yaml that belong to one machine: where it
// listens, its certificates, how it reaches the store, its own secrets and
// files. Everything else is the cluster's, kept in the store and the same on
// every control plane (see Layer).
var localKeys = [][]string{
	{"listen"}, {"admin_listen"}, {"metrics_listen"}, {"agent_listen"},
	{"max_connections"}, {"admin_token_file"}, {"admin_remote"}, {"debug_pprof"},
	{"agent_insecure"}, {"agent_tls"},
	{"auth", "store"}, {"threads", "store"},
	{"secrets"}, {"backup_dir"}, {"backup_copy_key"},
	{"tracing"}, {"monitoring"},
	{"mirror", "api"}, {"mirror", "tls"},
	{"endpoints"},
}

// IsLocal is whether a key path is one machine's own, kept in its
// server.yaml, rather than the cluster's.
func IsLocal(path []string) bool {
	for _, l := range localKeys {
		if len(path) >= len(l) && slices.Equal(path[:len(l)], l) {
			return true
		}
	}
	return false
}

// LocalKeys are the per-machine keys, dotted, for saying which they are.
func LocalKeys() []string {
	out := make([]string, len(localKeys))
	for i, l := range localKeys {
		out[i] = strings.Join(l, ".")
	}
	return out
}

// ReadFile is a configuration file as YAML maps, empty for no path.
func ReadFile(path string) (map[string]any, error) {
	if path == "" {
		return map[string]any{}, nil
	}
	b, err := os.ReadFile(path) //nolint:gosec // the configuration
	if err != nil {
		return nil, fmt.Errorf("config file: %w", err)
	}
	m, err := Parse(b)
	if err != nil {
		return nil, fmt.Errorf("parsing %s: %w", path, err)
	}
	return m, nil
}

// Parse reads a YAML document of settings.
func Parse(b []byte) (map[string]any, error) {
	m := map[string]any{}
	if len(bytes.TrimSpace(b)) == 0 {
		return m, nil
	}
	if err := yaml.Unmarshal(b, &m); err != nil {
		return nil, err
	}
	if m == nil {
		m = map[string]any{}
	}
	return m, nil
}

// Split separates settings into one machine's and the cluster's.
func Split(m map[string]any) (local, cluster map[string]any) {
	local, cluster = map[string]any{}, map[string]any{}
	split(m, nil, local, cluster)
	return local, cluster
}

func split(m map[string]any, path []string, local, cluster map[string]any) {
	for k, v := range m {
		p := append(slices.Clone(path), k)
		switch {
		case IsLocal(p):
			local[k] = v
		case hasLocalBelow(p):
			sub, ok := v.(map[string]any)
			if !ok {
				cluster[k] = v
				continue
			}
			l, c := map[string]any{}, map[string]any{}
			split(sub, p, l, c)
			if len(l) > 0 {
				local[k] = l
			}
			if len(c) > 0 {
				cluster[k] = c
			}
		default:
			cluster[k] = v
		}
	}
}

// hasLocalBelow is whether some local key lies under a path: auth holds
// store, which is local, beside its policies, which are not.
func hasLocalBelow(path []string) bool {
	for _, l := range localKeys {
		if len(l) > len(path) && slices.Equal(l[:len(path)], path) {
			return true
		}
	}
	return false
}

// Layer is what a control plane runs on: its own keys from its server.yaml
// and every other one from the cluster's settings. The file's cluster keys
// are not read once the cluster has settings of its own; Ignored names them.
func Layer(file, cluster map[string]any) map[string]any {
	local, _ := Split(file)
	out := clone(local)
	merge(out, cluster)
	return out
}

// Ignored are the cluster keys a file has that the cluster's settings, once
// there, take the place of.
func Ignored(file map[string]any) []string {
	_, c := Split(file)
	return flatKeys(c, "")
}

func flatKeys(m map[string]any, prefix string) []string {
	var out []string
	for _, k := range slices.Sorted(maps.Keys(m)) {
		p := prefix + k
		if sub, ok := m[k].(map[string]any); ok && hasLocalBelow(strings.Split(p, ".")) {
			out = append(out, flatKeys(sub, p+".")...)
			continue
		}
		out = append(out, p)
	}
	return out
}

// CheckCluster refuses a cluster document that sets one machine's keys: a
// listener or a certificate path written there would be every control
// plane's, and none of them could start on it.
func CheckCluster(m map[string]any) error {
	local, _ := Split(m)
	if keys := flatKeys(local, ""); len(keys) > 0 {
		return fmt.Errorf("%s belong to each machine's server.yaml, not to the cluster's settings",
			strings.Join(keys, ", "))
	}
	return nil
}

// Changed are the keys whose values differ between two sets of settings,
// at the top level, and below it where a key mixes local and cluster ones.
func Changed(a, b map[string]any) []string {
	var out []string
	changed(a, b, nil, &out)
	slices.Sort(out)
	return out
}

func changed(a, b map[string]any, path []string, out *[]string) {
	keys := map[string]bool{}
	for k := range a {
		keys[k] = true
	}
	for k := range b {
		keys[k] = true
	}
	for k := range keys {
		p := append(slices.Clone(path), k)
		va, vb := a[k], b[k]
		if reflect.DeepEqual(va, vb) {
			continue
		}
		ma, okA := va.(map[string]any)
		mb, okB := vb.(map[string]any)
		if hasLocalBelow(p) && (okA || va == nil) && (okB || vb == nil) {
			changed(ma, mb, p, out)
			continue
		}
		*out = append(*out, strings.Join(p, "."))
	}
}

// Marshal writes settings as YAML.
func Marshal(m map[string]any) ([]byte, error) {
	if len(m) == 0 {
		return nil, nil
	}
	var b bytes.Buffer
	enc := yaml.NewEncoder(&b)
	enc.SetIndent(2)
	if err := enc.Encode(m); err != nil {
		return nil, err
	}
	return b.Bytes(), enc.Close()
}

// FromMap is the configuration settings make: laid over the defaults, with
// LLMGW_ environment overrides, and validated.
func FromMap(m map[string]any) (Server, error) {
	b, err := Marshal(m)
	if err != nil {
		return Default(), err
	}
	return load(b)
}

// bytesProvider hands koanf a document already read.
type bytesProvider []byte

func (b bytesProvider) ReadBytes() ([]byte, error) { return b, nil }
func (b bytesProvider) Read() (map[string]any, error) {
	return nil, errors.New("bytesProvider reads bytes only")
}

func clone(m map[string]any) map[string]any {
	out := make(map[string]any, len(m))
	for k, v := range m {
		if sub, ok := v.(map[string]any); ok {
			v = clone(sub)
		}
		out[k] = v
	}
	return out
}

// merge lays src over dst: maps key by key, everything else replaced.
func merge(dst, src map[string]any) {
	for k, v := range src {
		if sv, ok := v.(map[string]any); ok {
			if dv, ok := dst[k].(map[string]any); ok {
				merge(dv, sv)
				continue
			}
			v = clone(sv)
		}
		dst[k] = v
	}
}