CrossRelay/internal/codexcfg/codexcfg.go

// Package codexcfg reads and writes Codex's config.toml. CrossRelay
// only touches the keys it manages; every other line (tables,
// comments, spacing, line endings) round-trips untouched.
package codexcfg

import (
	"bytes"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"
)

// ProviderID is the model_providers entry CrossRelay manages.
const ProviderID = "crossrelay"

// TableName is the full TOML table path of the managed provider.
const TableName = "model_providers." + ProviderID

// DefaultPath returns the location of Codex's config.toml. CODEX_HOME
// overrides the default location (~/.codex).
func DefaultPath() (string, error) {
	if v := strings.TrimSpace(os.Getenv("CODEX_HOME")); v != "" {
		return filepath.Join(v, "config.toml"), nil
	}
	home, err := os.UserHomeDir()
	if err != nil {
		return "", fmt.Errorf("locate home directory: %w", err)
	}
	return filepath.Join(home, ".codex", "config.toml"), nil
}

// Manager performs guarded edits on one config.toml file.
type Manager struct {
	path      string
	backupDir string
	mu        sync.Mutex
}

// New builds a manager for configPath, storing backups in backupDir.
func New(configPath, backupDir string) *Manager {
	return &Manager{path: configPath, backupDir: backupDir}
}

// ApplyOptions describes what gets injected into config.toml.
type ApplyOptions struct {
	BaseURL string
	Token   string
}

// Injection is the managed state written into config.toml. It doubles
// as the preview returned to the UI after an apply.
type Injection struct {
	ModelProvider string            `json:"model_provider"`
	Table         map[string]string `json:"table"`
	// PriorProvider is the model_provider value found before the
	// apply, empty when the file carried none.
	PriorProvider string `json:"prior_provider,omitempty"`
}

// Apply injects the relay configuration, backing up the current file
// first when the content changes. It returns the resulting injection.
func (m *Manager) Apply(opts ApplyOptions) (*Injection, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	original, err := m.readRaw()
	if err != nil {
		return nil, err
	}
	lines := splitLines(original)
	prior := readTopLevelKey(lines, "model_provider")
	lines = setTopLevelKey(lines, "model_provider", ProviderID)
	lines = setTable(lines, TableName, managedTableLines(opts))
	updated := joinLines(lines)
	if original == nil || !bytes.Equal(original, updated) {
		if len(original) > 0 {
			if err := m.backupLocked(original); err != nil {
				return nil, err
			}
		}
		if err := m.writeFile(updated); err != nil {
			return nil, err
		}
	}
	return &Injection{
		ModelProvider: ProviderID,
		Table: map[string]string{
			"name":                      ProviderID,
			"wire_api":                  "responses",
			"base_url":                  opts.BaseURL,
			"experimental_bearer_token": opts.Token,
			"requires_openai_auth":      "false",
		},
		PriorProvider: prior,
	}, nil
}

// RemoveManaged deletes the keys CrossRelay manages. priorProvider is
// written back into model_provider when non-empty; otherwise the key
// is dropped, and only while it still carries the managed value.
func (m *Manager) RemoveManaged(priorProvider string) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	original, err := m.readRaw()
	if err != nil {
		return err
	}
	if len(original) == 0 {
		return nil
	}
	lines := splitLines(original)
	lines = removeTable(lines, TableName)
	if prior := strings.TrimSpace(priorProvider); prior != "" {
		lines = setTopLevelKey(lines, "model_provider", prior)
	} else {
		lines = removeTopLevelKey(lines, "model_provider", ProviderID)
	}
	updated := joinLines(lines)
	if bytes.Equal(original, updated) {
		return nil
	}
	return m.writeFile(updated)
}

// RestoreBackup copies a backup file over config.toml and returns the
// backup name.
func (m *Manager) RestoreBackup(name string) (string, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if !strings.HasPrefix(name, "config-") || !strings.HasSuffix(name, ".toml") ||
		strings.ContainsAny(name, `/\`) {
		return "", fmt.Errorf("invalid backup name %q", name)
	}
	data, err := os.ReadFile(filepath.Join(m.backupDir, name))
	if err != nil {
		return "", err
	}
	if err := m.writeFile(data); err != nil {
		return "", err
	}
	return name, nil
}

// BackupInfo describes one backup file.
type BackupInfo struct {
	Name     string    `json:"name"`
	Size     int64     `json:"size"`
	Modified time.Time `json:"modified"`
}

// Backups lists backup files, newest first.
func (m *Manager) Backups() ([]BackupInfo, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.listBackupsLocked()
}

func managedTableLines(opts ApplyOptions) []string {
	return []string{
		fmt.Sprintf("name = %q", ProviderID),
		`wire_api = "responses"`,
		fmt.Sprintf("base_url = %q", opts.BaseURL),
		fmt.Sprintf("experimental_bearer_token = %q", opts.Token),
		"requires_openai_auth = false",
	}
}

func (m *Manager) readRaw() ([]byte, error) {
	data, err := os.ReadFile(m.path)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, nil
		}
		return nil, err
	}
	return data, nil
}

func (m *Manager) writeFile(data []byte) error {
	if err := os.MkdirAll(filepath.Dir(m.path), 0o700); err != nil {
		return err
	}
	tmp := m.path + ".tmp"
	if err := os.WriteFile(tmp, data, 0o600); err != nil {
		return err
	}
	if err := os.Rename(tmp, m.path); err != nil {
		_ = os.Remove(tmp)
		return err
	}
	return nil
}

func (m *Manager) backupLocked(data []byte) error {
	if err := os.MkdirAll(m.backupDir, 0o700); err != nil {
		return err
	}
	name := fmt.Sprintf("config-%s.toml", time.Now().Format("20060102-150405"))
	if err := os.WriteFile(filepath.Join(m.backupDir, name), data, 0o600); err != nil {
		return err
	}
	m.pruneLocked(20)
	return nil
}

func (m *Manager) listBackupsLocked() ([]BackupInfo, error) {
	entries, err := os.ReadDir(m.backupDir)
	if err != nil {
		if os.IsNotExist(err) {
			return []BackupInfo{}, nil
		}
		return nil, err
	}
	out := []BackupInfo{}
	for _, e := range entries {
		if e.IsDir() || !strings.HasPrefix(e.Name(), "config-") || !strings.HasSuffix(e.Name(), ".toml") {
			continue
		}
		info, err := e.Info()
		if err != nil {
			continue
		}
		out = append(out, BackupInfo{Name: e.Name(), Size: info.Size(), Modified: info.ModTime()})
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Modified.After(out[j].Modified) })
	return out, nil
}

func (m *Manager) pruneLocked(keep int) {
	backups, err := m.listBackupsLocked()
	if err != nil || len(backups) <= keep {
		return
	}
	for _, b := range backups[keep:] {
		_ = os.Remove(filepath.Join(m.backupDir, b.Name))
	}
}

// splitLines splits file content on "\n". A trailing newline shows up
// as a final empty element, so joining reproduces the file exactly,
// CRLF endings included.
func splitLines(data []byte) []string {
	if len(data) == 0 {
		return nil
	}
	return strings.Split(string(data), "\n")
}

func joinLines(lines []string) []byte {
	out := strings.Join(lines, "\n")
	if len(lines) > 0 && lines[len(lines)-1] != "" {
		out += "\n"
	}
	return []byte(out)
}

func isHeaderLine(line string) bool {
	return strings.HasPrefix(strings.TrimSpace(line), "[")
}

func parseKeyValue(line string) (key, value string, ok bool) {
	s := strings.TrimSpace(line)
	if s == "" || strings.HasPrefix(s, "#") || strings.HasPrefix(s, "[") {
		return "", "", false
	}
	eq := strings.Index(s, "=")
	if eq < 0 {
		return "", "", false
	}
	key = strings.Trim(strings.TrimSpace(s[:eq]), `"'`)
	if key == "" {
		return "", "", false
	}
	return key, unquoteTOML(strings.TrimSpace(s[eq+1:])), true
}

func unquoteTOML(v string) string {
	if len(v) >= 2 && v[0] == '"' {
		if end := strings.Index(v[1:], `"`); end >= 0 {
			if uq, err := strconv.Unquote(v[:end+2]); err == nil {
				return uq
			}
			return v[1 : end+1]
		}
	}
	if len(v) >= 2 && v[0] == '\'' {
		if end := strings.Index(v[1:], "'"); end >= 0 {
			return v[1 : end+1]
		}
	}
	if i := strings.Index(v, "#"); i >= 0 {
		v = strings.TrimSpace(v[:i])
	}
	return v
}

// findTopLevelKey returns the index of a key assignment in the
// top-level section (before the first table header).
func findTopLevelKey(lines []string, key string) int {
	for i, line := range lines {
		if isHeaderLine(line) {
			return -1
		}
		if k, _, ok := parseKeyValue(line); ok && k == key {
			return i
		}
	}
	return -1
}

func readTopLevelKey(lines []string, key string) string {
	if i := findTopLevelKey(lines, key); i >= 0 {
		if _, v, ok := parseKeyValue(lines[i]); ok {
			return v
		}
	}
	return ""
}

// setTopLevelKey replaces an existing top-level key or inserts one at
// the end of the top-level section.
func setTopLevelKey(lines []string, key, value string) []string {
	line := fmt.Sprintf("%s = %q", key, value)
	if i := findTopLevelKey(lines, key); i >= 0 {
		out := append([]string{}, lines...)
		out[i] = line
		return out
	}
	return insertTopLevelLine(lines, line)
}

func insertTopLevelLine(lines []string, line string) []string {
	idx := len(lines)
	for i, l := range lines {
		if isHeaderLine(l) {
			idx = i
			break
		}
	}
	// A trailing empty element represents the file's final newline; new
	// lines go before it.
	if idx == len(lines) && idx > 0 && lines[idx-1] == "" {
		idx--
	}
	out := make([]string, 0, len(lines)+1)
	out = append(out, lines[:idx]...)
	out = append(out, line)
	out = append(out, lines[idx:]...)
	return out
}

// removeTopLevelKey drops a top-level key line, but only while it
// still carries the given value.
func removeTopLevelKey(lines []string, key, value string) []string {
	i := findTopLevelKey(lines, key)
	if i < 0 {
		return lines
	}
	_, v, ok := parseKeyValue(lines[i])
	if !ok || v != value {
		return lines
	}
	out := make([]string, 0, len(lines))
	out = append(out, lines[:i]...)
	out = append(out, lines[i+1:]...)
	return out
}

// setTable replaces the body of an existing table or appends a new
// table at the end of the file.
func setTable(lines []string, name string, body []string) []string {
	header := "[" + name + "]"
	if i := findTable(lines, name); i >= 0 {
		end := tableEnd(lines, i)
		var sep []string
		if end < len(lines) && isHeaderLine(lines[end]) && end > i+1 && strings.TrimSpace(lines[end-1]) == "" {
			sep = []string{""}
		}
		out := make([]string, 0, len(lines)+len(body)+1)
		out = append(out, lines[:i+1]...)
		out = append(out, body...)
		out = append(out, sep...)
		out = append(out, lines[end:]...)
		return out
	}
	block := append([]string{header}, body...)
	end := len(lines)
	if end > 0 && lines[end-1] == "" {
		end--
	}
	if end > 0 && strings.TrimSpace(lines[end-1]) != "" {
		block = append([]string{""}, block...)
	}
	out := make([]string, 0, len(lines)+len(block)+1)
	out = append(out, lines[:end]...)
	out = append(out, block...)
	out = append(out, lines[end:]...)
	return out
}

// removeTable drops the named table plus one preceding blank
// separator so no double blank lines are left behind.
func removeTable(lines []string, name string) []string {
	i := findTable(lines, name)
	if i < 0 {
		return lines
	}
	end := tableEnd(lines, i)
	start := i
	if start > 0 && strings.TrimSpace(lines[start-1]) == "" {
		after := end >= len(lines) || strings.TrimSpace(lines[end]) == "" || isHeaderLine(lines[end])
		if after {
			start--
		}
	}
	out := make([]string, 0, len(lines))
	out = append(out, lines[:start]...)
	out = append(out, lines[end:]...)
	return out
}

func findTable(lines []string, name string) int {
	want := "[" + name + "]"
	for i, line := range lines {
		if strings.TrimSpace(line) == want {
			return i
		}
	}
	return -1
}

func tableEnd(lines []string, headerIdx int) int {
	end := headerIdx + 1
	for end < len(lines) && !isHeaderLine(lines[end]) {
		end++
	}
	return end
}