config.go

// Copyright (c) Tailscale Inc & AUTHORS
// SPDX-License-Identifier: BSD-3-Clause

package main

import (
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"sync/atomic"

	"tailscale.com/atomicfile"
)

// CompressSettings are the runtime-tunable knobs forwarded to headroom.compress.
// They mirror headroom's CompressConfig field names exactly (the worker passes
// them straight through as kwargs), and default to headroom's own defaults so
// behavior is unchanged until someone tunes them.
//
// TargetRatio and KompressModel are pointers because their "unset" value is
// JSON null with real meaning (target_ratio null = "model decides ~15%";
// kompress_model null = headroom's default model). SavingsProfile is likewise a
// pointer (null = no profile).
//
// Deliberately NOT mirrored: CompressConfig.frozen_message_count (added in
// headroom 0.32.0). It is not a static policy knob — it's a per-request value
// ("the message count of the previous request") that a caller managing its own
// conversation loop passes so compression won't rewrite the already-cached
// prefix. A single process-global value here would be meaningless (or harmful if
// nonzero). Setting it correctly requires per-session prior-turn state, i.e. the
// stateful/delta-worker rework tracked in TODO.md; until then we leave it at
// headroom's default (0 = no frozen prefix), which is the pre-0.32 behavior.
type CompressSettings struct {
	CompressUserMessages   bool     `json:"compress_user_messages"`
	CompressSystemMessages bool     `json:"compress_system_messages"`
	ProtectRecent          int      `json:"protect_recent"`
	ProtectAnalysisContext bool     `json:"protect_analysis_context"`
	TargetRatio            *float64 `json:"target_ratio"`
	MinTokensToCompress    int      `json:"min_tokens_to_compress"`
	KompressModel          *string  `json:"kompress_model"`

	// SavingsProfile selects a headroom named profile (>= 0.26.0). When set it
	// OVERRIDES the individual knobs above. null = no profile (default).
	SavingsProfile *string `json:"savings_profile"`
}

// savingsProfiles is the set of named profiles headroom exposes via the library
// compress() path. Validation rejects anything else so an unknown name can never
// reach compress() — where it would raise (it's applied before compress()'s own
// fail-open try) and make the pool recycle the worker on every request.
//
// agent-90 and balanced date to 0.26.0; the coding and general workload personas
// were added in 0.30.0. All four apply cleanly via compress() (the library path
// sets only the standard CompressConfig knobs each profile carries). Keep this in
// sync with headroom.agent_savings._PROFILES.
var savingsProfiles = map[string]bool{
	"agent-90": true,
	"balanced": true,
	"coding":   true,
	"general":  true,
}

// defaultSettings mirrors headroom's CompressConfig defaults (compress.py).
func defaultSettings() CompressSettings {
	return CompressSettings{
		CompressUserMessages:   false,
		CompressSystemMessages: true,
		ProtectRecent:          4,
		ProtectAnalysisContext: true,
		TargetRatio:            nil,
		MinTokensToCompress:    250,
		KompressModel:          nil,
	}
}

// validate rejects nonsensical settings so a bad PUT can't wedge the service.
func (s CompressSettings) validate() error {
	if s.ProtectRecent < 0 {
		return fmt.Errorf("protect_recent must be >= 0")
	}
	if s.MinTokensToCompress < 0 {
		return fmt.Errorf("min_tokens_to_compress must be >= 0")
	}
	if s.TargetRatio != nil && (*s.TargetRatio <= 0 || *s.TargetRatio > 1) {
		return fmt.Errorf("target_ratio must be in (0, 1] or null")
	}
	if s.KompressModel != nil && *s.KompressModel == "" {
		return fmt.Errorf("kompress_model must be a non-empty string or null")
	}
	if s.SavingsProfile != nil && !savingsProfiles[*s.SavingsProfile] {
		return fmt.Errorf("savings_profile must be one of %s, or null", strings.Join(sortedProfiles(), ", "))
	}
	return nil
}

// sortedProfiles returns the accepted savings_profile names in stable order, so
// the validation error text is deterministic and always matches savingsProfiles.
func sortedProfiles() []string {
	names := make([]string, 0, len(savingsProfiles))
	for name := range savingsProfiles {
		names = append(names, name)
	}
	sort.Strings(names)
	return names
}

// textEnabled reports whether a knob that drives the ML text compressor
// (Kompress) is on. It's the single source of truth for whether workers should
// preload that model at startup (passed to them via TSHEADROOM_PRELOAD; see
// worker.py's _warmup). Keep this in sync with any new ML-driving knob added to
// CompressSettings.
//
// CompressSystemMessages defaults to true, so with headroom-ai[ml] installed the
// ML model is needed for ordinary traffic even under an otherwise-default config
// (and headroom also uses Kompress as its fallback for tool/mixed content). We
// therefore include it here, which makes preload on by default — workers come up
// warm instead of cold-loading the ~600MB model on their first live request.
func (s CompressSettings) textEnabled() bool {
	return s.CompressUserMessages || s.CompressSystemMessages || s.TargetRatio != nil
}

// settingsStore holds the current settings behind an atomic pointer (read once
// per request, swapped on update) and persists them to disk so the service
// starts up with the last state.
type settingsStore struct {
	cur  atomic.Pointer[CompressSettings]
	path string // "" disables persistence
	log  *slog.Logger

	mu sync.Mutex // serializes writers (the full read-merge-validate-save-swap)
}

// loadSettings builds a store, seeding it from path if present and valid,
// otherwise from defaults. A missing or corrupt file is non-fatal.
func loadSettings(path string, log *slog.Logger) *settingsStore {
	st := &settingsStore{path: path, log: log}
	s := defaultSettings()

	if path != "" {
		switch b, err := os.ReadFile(path); {
		case err == nil:
			var loaded CompressSettings
			if jerr := json.Unmarshal(b, &loaded); jerr != nil {
				log.Warn("config file unparseable; using defaults", "path", path, "err", jerr)
			} else if verr := loaded.validate(); verr != nil {
				log.Warn("config file invalid; using defaults", "path", path, "err", verr)
			} else {
				s = loaded
				log.Info("loaded config", "path", path)
			}
		case os.IsNotExist(err):
			log.Info("no config file yet; using defaults", "path", path)
		default:
			log.Warn("config file unreadable; using defaults", "path", path, "err", err)
		}
	}

	st.cur.Store(&s)
	return st
}

// get returns a snapshot of the current settings.
func (st *settingsStore) get() CompressSettings { return *st.cur.Load() }

// set validates, persists, then swaps in a full settings value. On a
// persistence error the in-memory settings are left unchanged so disk and
// memory stay consistent.
func (st *settingsStore) set(s CompressSettings) error {
	st.mu.Lock()
	defer st.mu.Unlock()
	return st.applyLocked(s)
}

// merge applies a partial JSON update onto the current settings as one atomic
// read-modify-write, so concurrent PUTs can't lose each other's changes.
// Returns the resulting settings.
func (st *settingsStore) merge(body []byte) (CompressSettings, error) {
	st.mu.Lock()
	defer st.mu.Unlock()

	merged := *st.cur.Load()
	// Unmarshaling onto the current value leaves omitted fields untouched
	// (partial update) and lets explicit null clear a pointer field.
	if err := json.Unmarshal(body, &merged); err != nil {
		return CompressSettings{}, err
	}
	if err := st.applyLocked(merged); err != nil {
		return CompressSettings{}, err
	}
	return merged, nil
}

// applyLocked validates, persists, and swaps in s. The caller must hold st.mu.
func (st *settingsStore) applyLocked(s CompressSettings) error {
	if err := s.validate(); err != nil {
		return err
	}
	if err := st.save(s); err != nil {
		return fmt.Errorf("persist config: %w", err)
	}
	st.cur.Store(&s)
	return nil
}

// save atomically writes settings to disk. The caller must hold st.mu.
func (st *settingsStore) save(s CompressSettings) error {
	if st.path == "" {
		return nil
	}
	b, err := json.MarshalIndent(s, "", "  ")
	if err != nil {
		return err
	}
	if dir := filepath.Dir(st.path); dir != "" {
		if err := os.MkdirAll(dir, 0o755); err != nil {
			return err
		}
	}
	// atomicfile.WriteFile writes to a unique temp file in the same dir, fsyncs,
	// and renames into place (cleaning up the temp on error).
	return atomicfile.WriteFile(st.path, append(b, '\n'), 0o644)
}

// configHandler serves the runtime tuning API: GET returns the current
// settings (plus the detected headroom_version); PUT merges the provided fields
// onto the current settings (partial updates allowed), validates, persists, and
// returns the result.
type configHandler struct {
	store *settingsStore
	log   *slog.Logger

	// headroomVersion reports the detected headroom-ai version and whether it's
	// known yet. Used to surface the version on GET and to reject setting
	// savings_profile on a version that predates it. nil in tests = unknown.
	headroomVersion func() (string, bool)
}

// configView is the GET response: the settings plus the read-only detected
// headroom version. Embedding flattens CompressSettings' fields, so
// savings_profile stays visible even when the running headroom can't honor it.
type configView struct {
	CompressSettings
	HeadroomVersion string `json:"headroom_version,omitempty"`
}

const maxConfigBody = 64 << 10

func (h *configHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	switch r.Method {
	case http.MethodGet:
		view := configView{CompressSettings: h.store.get()}
		if h.headroomVersion != nil {
			if ver, ok := h.headroomVersion(); ok {
				view.HeadroomVersion = ver
			}
		}
		writeJSON(w, http.StatusOK, view)

	case http.MethodPut, http.MethodPost:
		body, err := io.ReadAll(io.LimitReader(r.Body, maxConfigBody))
		if err != nil {
			http.Error(w, "read body failed", http.StatusBadRequest)
			return
		}
		// Reject *setting* savings_profile on a headroom too old to honor it
		// (where it would be a silent no-op), before we validate/persist.
		// Clearing it (null) or leaving it untouched is always allowed.
		if h.settingNonNullSavingsProfile(body) {
			if ver, ok := h.detectedVersion(); ok && !supportsSavingsProfile(ver) {
				http.Error(w, fmt.Sprintf("savings_profile requires headroom-ai >= 0.26.0 (detected %s)", ver), http.StatusBadRequest)
				return
			}
		}
		updated, err := h.store.merge(body)
		if err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}
		h.log.Info("config updated", "settings", updated)
		writeJSON(w, http.StatusOK, updated)

	default:
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
	}
}

// detectedVersion returns the detected headroom version and whether it's known,
// nil-safe for tests that don't wire the accessor.
func (h *configHandler) detectedVersion() (string, bool) {
	if h.headroomVersion == nil {
		return "", false
	}
	return h.headroomVersion()
}

// settingNonNullSavingsProfile reports whether the PUT body assigns
// savings_profile a non-null value (i.e. is trying to turn it on). Absent key or
// explicit null is not "setting" it.
func (h *configHandler) settingNonNullSavingsProfile(body []byte) bool {
	var probe map[string]json.RawMessage
	if err := json.Unmarshal(body, &probe); err != nil {
		return false // malformed body; merge will reject it with a clearer error
	}
	raw, ok := probe["savings_profile"]
	return ok && strings.TrimSpace(string(raw)) != "null"
}

// writeJSON writes v as a JSON response with the given status.
func writeJSON(w http.ResponseWriter, status int, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(v)
}