package config import ( "bytes" "fmt" "os" "strings" "gopkg.in/yaml.v3" ) // SaveConfigPreserveComments writes the config back to YAML while preserving existing comments // and key ordering by loading the original file into a yaml.Node tree and updating values in-place. func SaveConfigPreserveComments(configFile string, cfg *Config) error { persistCfg := sanitizeConfigForPersist(cfg) // Load original YAML as a node tree to preserve comments and ordering. data, err := os.ReadFile(configFile) if err != nil { return err } var original yaml.Node if err = yaml.Unmarshal(data, &original); err != nil { return err } if original.Kind != yaml.DocumentNode || len(original.Content) == 0 { return fmt.Errorf("invalid yaml document structure") } if original.Content[0] == nil || original.Content[0].Kind != yaml.MappingNode { return fmt.Errorf("expected root mapping node") } // Marshal the current cfg to YAML, then unmarshal to a yaml.Node we can merge from. rendered, err := yaml.Marshal(persistCfg) if err != nil { return err } var generated yaml.Node if err = yaml.Unmarshal(rendered, &generated); err != nil { return err } if generated.Kind != yaml.DocumentNode || len(generated.Content) == 0 || generated.Content[0] == nil { return fmt.Errorf("invalid generated yaml structure") } if generated.Content[0].Kind != yaml.MappingNode { return fmt.Errorf("expected generated root mapping node") } // Remove deprecated sections before merging back the sanitized config. removeLegacyAuthBlock(original.Content[0]) removeLegacyOpenAICompatAPIKeys(original.Content[0]) removeLegacyAmpKeys(original.Content[0]) removeLegacyGenerativeLanguageKeys(original.Content[0]) pruneMappingToGeneratedKeys(original.Content[0], generated.Content[0], "oauth-excluded-models") // Merge generated into original in-place, preserving comments/order of existing nodes. mergeMappingPreserve(original.Content[0], generated.Content[0]) normalizeCollectionNodeStyles(original.Content[0]) // Write back. f, err := os.Create(configFile) if err != nil { return err } defer func() { _ = f.Close() }() var buf bytes.Buffer enc := yaml.NewEncoder(&buf) enc.SetIndent(2) if err = enc.Encode(&original); err != nil { _ = enc.Close() return err } if err = enc.Close(); err != nil { return err } data = NormalizeCommentIndentation(buf.Bytes()) _, err = f.Write(data) return err } // SaveConfigPreserveCommentsUpdateNestedScalar updates a nested scalar key path like ["a","b"] // while preserving comments and positions. func SaveConfigPreserveCommentsUpdateNestedScalar(configFile string, path []string, value string) error { data, err := os.ReadFile(configFile) if err != nil { return err } var root yaml.Node if err = yaml.Unmarshal(data, &root); err != nil { return err } if root.Kind != yaml.DocumentNode || len(root.Content) == 0 { return fmt.Errorf("invalid yaml document structure") } node := root.Content[0] // descend mapping nodes following path for i, key := range path { if i == len(path)-1 { // set final scalar v := getOrCreateMapValue(node, key) v.Kind = yaml.ScalarNode v.Tag = "!!str" v.Value = value } else { next := getOrCreateMapValue(node, key) if next.Kind != yaml.MappingNode { next.Kind = yaml.MappingNode next.Tag = "!!map" } node = next } } f, err := os.Create(configFile) if err != nil { return err } defer func() { _ = f.Close() }() var buf bytes.Buffer enc := yaml.NewEncoder(&buf) enc.SetIndent(2) if err = enc.Encode(&root); err != nil { _ = enc.Close() return err } if err = enc.Close(); err != nil { return err } data = NormalizeCommentIndentation(buf.Bytes()) _, err = f.Write(data) return err } // NormalizeCommentIndentation removes indentation from standalone YAML comment lines to keep them left aligned. func NormalizeCommentIndentation(data []byte) []byte { lines := bytes.Split(data, []byte("\n")) changed := false for i, line := range lines { trimmed := bytes.TrimLeft(line, " \t") if len(trimmed) == 0 || trimmed[0] != '#' { continue } if len(trimmed) == len(line) { continue } lines[i] = append([]byte(nil), trimmed...) changed = true } if !changed { return data } return bytes.Join(lines, []byte("\n")) } // getOrCreateMapValue finds the value node for a given key in a mapping node. // If not found, it appends a new key/value pair and returns the new value node. func getOrCreateMapValue(mapNode *yaml.Node, key string) *yaml.Node { if mapNode.Kind != yaml.MappingNode { mapNode.Kind = yaml.MappingNode mapNode.Tag = "!!map" mapNode.Content = nil } for i := 0; i+1 < len(mapNode.Content); i += 2 { k := mapNode.Content[i] if k.Value == key { return mapNode.Content[i+1] } } // append new key/value mapNode.Content = append(mapNode.Content, &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: key}) val := &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: ""} mapNode.Content = append(mapNode.Content, val) return val } // mergeMappingPreserve merges keys from src into dst mapping node while preserving // key order and comments of existing keys in dst. New keys are only added if their // value is non-zero to avoid polluting the config with defaults. func mergeMappingPreserve(dst, src *yaml.Node) { if dst == nil || src == nil { return } if dst.Kind != yaml.MappingNode || src.Kind != yaml.MappingNode { // If kinds do not match, prefer replacing dst with src semantics in-place // but keep dst node object to preserve any attached comments at the parent level. copyNodeShallow(dst, src) return } for i := 0; i+1 < len(src.Content); i += 2 { sk := src.Content[i] sv := src.Content[i+1] idx := findMapKeyIndex(dst, sk.Value) if idx >= 0 { // Merge into existing value node (always update, even to zero values) dv := dst.Content[idx+1] mergeNodePreserve(dv, sv) } else { // New key: only add if value is non-zero to avoid polluting config with defaults if isZeroValueNode(sv) { continue } dst.Content = append(dst.Content, deepCopyNode(sk), deepCopyNode(sv)) } } } // mergeNodePreserve merges src into dst for scalars, mappings and sequences while // reusing destination nodes to keep comments and anchors. For sequences, it updates // in-place by index. func mergeNodePreserve(dst, src *yaml.Node) { if dst == nil || src == nil { return } switch src.Kind { case yaml.MappingNode: if dst.Kind != yaml.MappingNode { copyNodeShallow(dst, src) } mergeMappingPreserve(dst, src) case yaml.SequenceNode: // Preserve explicit null style if dst was null and src is empty sequence if dst.Kind == yaml.ScalarNode && dst.Tag == "!!null" && len(src.Content) == 0 { // Keep as null to preserve original style return } if dst.Kind != yaml.SequenceNode { dst.Kind = yaml.SequenceNode dst.Tag = "!!seq" dst.Content = nil } reorderSequenceForMerge(dst, src) // Update elements in place minContent := len(dst.Content) if len(src.Content) < minContent { minContent = len(src.Content) } for i := 0; i < minContent; i++ { if dst.Content[i] == nil { dst.Content[i] = deepCopyNode(src.Content[i]) continue } mergeNodePreserve(dst.Content[i], src.Content[i]) if dst.Content[i] != nil && src.Content[i] != nil && dst.Content[i].Kind == yaml.MappingNode && src.Content[i].Kind == yaml.MappingNode { pruneMissingMapKeys(dst.Content[i], src.Content[i]) } } // Append any extra items from src for i := len(dst.Content); i < len(src.Content); i++ { dst.Content = append(dst.Content, deepCopyNode(src.Content[i])) } // Truncate if dst has extra items not in src if len(src.Content) < len(dst.Content) { dst.Content = dst.Content[:len(src.Content)] } case yaml.ScalarNode, yaml.AliasNode: // For scalars, update Tag and Value but keep Style from dst to preserve quoting dst.Kind = src.Kind dst.Tag = src.Tag dst.Value = src.Value // Keep dst.Style as-is intentionally case 0: // Unknown/empty kind; do nothing default: // Fallback: replace shallowly copyNodeShallow(dst, src) } } // findMapKeyIndex returns the index of key node in dst mapping (index of key, not value). // Returns -1 when not found. func findMapKeyIndex(mapNode *yaml.Node, key string) int { if mapNode == nil || mapNode.Kind != yaml.MappingNode { return -1 } for i := 0; i+1 < len(mapNode.Content); i += 2 { if mapNode.Content[i] != nil && mapNode.Content[i].Value == key { return i } } return -1 } // isZeroValueNode returns true if the YAML node represents a zero/default value // that should not be written as a new key to preserve config cleanliness. // For mappings and sequences, recursively checks if all children are zero values. func isZeroValueNode(node *yaml.Node) bool { if node == nil { return true } switch node.Kind { case yaml.ScalarNode: switch node.Tag { case "!!bool": return node.Value == "false" case "!!int", "!!float": return node.Value == "0" || node.Value == "0.0" case "!!str": return node.Value == "" case "!!null": return true } case yaml.SequenceNode: if len(node.Content) == 0 { return true } // Check if all elements are zero values for _, child := range node.Content { if !isZeroValueNode(child) { return false } } return true case yaml.MappingNode: if len(node.Content) == 0 { return true } // Check if all values are zero values (values are at odd indices) for i := 1; i < len(node.Content); i += 2 { if !isZeroValueNode(node.Content[i]) { return false } } return true } return false } // deepCopyNode creates a deep copy of a yaml.Node graph. func deepCopyNode(n *yaml.Node) *yaml.Node { if n == nil { return nil } cp := *n if len(n.Content) > 0 { cp.Content = make([]*yaml.Node, len(n.Content)) for i := range n.Content { cp.Content[i] = deepCopyNode(n.Content[i]) } } return &cp } // copyNodeShallow copies type/tag/value and resets content to match src, but // keeps the same destination node pointer to preserve parent relations/comments. func copyNodeShallow(dst, src *yaml.Node) { if dst == nil || src == nil { return } dst.Kind = src.Kind dst.Tag = src.Tag dst.Value = src.Value // Replace content with deep copy from src if len(src.Content) > 0 { dst.Content = make([]*yaml.Node, len(src.Content)) for i := range src.Content { dst.Content[i] = deepCopyNode(src.Content[i]) } } else { dst.Content = nil } } func reorderSequenceForMerge(dst, src *yaml.Node) { if dst == nil || src == nil { return } if len(dst.Content) == 0 { return } if len(src.Content) == 0 { return } original := append([]*yaml.Node(nil), dst.Content...) used := make([]bool, len(original)) ordered := make([]*yaml.Node, len(src.Content)) for i := range src.Content { if idx := matchSequenceElement(original, used, src.Content[i]); idx >= 0 { ordered[i] = original[idx] used[idx] = true } } dst.Content = ordered } func matchSequenceElement(original []*yaml.Node, used []bool, target *yaml.Node) int { if target == nil { return -1 } switch target.Kind { case yaml.MappingNode: id := sequenceElementIdentity(target) if id != "" { for i := range original { if used[i] || original[i] == nil || original[i].Kind != yaml.MappingNode { continue } if sequenceElementIdentity(original[i]) == id { return i } } } case yaml.ScalarNode: val := strings.TrimSpace(target.Value) if val != "" { for i := range original { if used[i] || original[i] == nil || original[i].Kind != yaml.ScalarNode { continue } if strings.TrimSpace(original[i].Value) == val { return i } } } default: } // Fallback to structural equality to preserve nodes lacking explicit identifiers. for i := range original { if used[i] || original[i] == nil { continue } if nodesStructurallyEqual(original[i], target) { return i } } return -1 } func sequenceElementIdentity(node *yaml.Node) string { if node == nil || node.Kind != yaml.MappingNode { return "" } identityKeys := []string{"id", "name", "alias", "api-key", "api_key", "apikey", "key", "provider", "model"} for _, k := range identityKeys { if v := mappingScalarValue(node, k); v != "" { return k + "=" + v } } for i := 0; i+1 < len(node.Content); i += 2 { keyNode := node.Content[i] valNode := node.Content[i+1] if keyNode == nil || valNode == nil || valNode.Kind != yaml.ScalarNode { continue } val := strings.TrimSpace(valNode.Value) if val != "" { return strings.ToLower(strings.TrimSpace(keyNode.Value)) + "=" + val } } return "" } func mappingScalarValue(node *yaml.Node, key string) string { if node == nil || node.Kind != yaml.MappingNode { return "" } lowerKey := strings.ToLower(key) for i := 0; i+1 < len(node.Content); i += 2 { keyNode := node.Content[i] valNode := node.Content[i+1] if keyNode == nil || valNode == nil || valNode.Kind != yaml.ScalarNode { continue } if strings.ToLower(strings.TrimSpace(keyNode.Value)) == lowerKey { return strings.TrimSpace(valNode.Value) } } return "" } func nodesStructurallyEqual(a, b *yaml.Node) bool { if a == nil || b == nil { return a == b } if a.Kind != b.Kind { return false } switch a.Kind { case yaml.MappingNode: if len(a.Content) != len(b.Content) { return false } for i := 0; i+1 < len(a.Content); i += 2 { if !nodesStructurallyEqual(a.Content[i], b.Content[i]) { return false } if !nodesStructurallyEqual(a.Content[i+1], b.Content[i+1]) { return false } } return true case yaml.SequenceNode: if len(a.Content) != len(b.Content) { return false } for i := range a.Content { if !nodesStructurallyEqual(a.Content[i], b.Content[i]) { return false } } return true case yaml.ScalarNode: return strings.TrimSpace(a.Value) == strings.TrimSpace(b.Value) case yaml.AliasNode: return nodesStructurallyEqual(a.Alias, b.Alias) default: return strings.TrimSpace(a.Value) == strings.TrimSpace(b.Value) } } func removeMapKey(mapNode *yaml.Node, key string) { if mapNode == nil || mapNode.Kind != yaml.MappingNode || key == "" { return } for i := 0; i+1 < len(mapNode.Content); i += 2 { if mapNode.Content[i] != nil && mapNode.Content[i].Value == key { mapNode.Content = append(mapNode.Content[:i], mapNode.Content[i+2:]...) return } } } func pruneMappingToGeneratedKeys(dstRoot, srcRoot *yaml.Node, key string) { if key == "" || dstRoot == nil || srcRoot == nil { return } if dstRoot.Kind != yaml.MappingNode || srcRoot.Kind != yaml.MappingNode { return } dstIdx := findMapKeyIndex(dstRoot, key) if dstIdx < 0 || dstIdx+1 >= len(dstRoot.Content) { return } srcIdx := findMapKeyIndex(srcRoot, key) if srcIdx < 0 { removeMapKey(dstRoot, key) return } if srcIdx+1 >= len(srcRoot.Content) { return } srcVal := srcRoot.Content[srcIdx+1] dstVal := dstRoot.Content[dstIdx+1] if srcVal == nil { dstRoot.Content[dstIdx+1] = nil return } if srcVal.Kind != yaml.MappingNode { dstRoot.Content[dstIdx+1] = deepCopyNode(srcVal) return } if dstVal == nil || dstVal.Kind != yaml.MappingNode { dstRoot.Content[dstIdx+1] = deepCopyNode(srcVal) return } pruneMissingMapKeys(dstVal, srcVal) } func pruneMissingMapKeys(dstMap, srcMap *yaml.Node) { if dstMap == nil || srcMap == nil || dstMap.Kind != yaml.MappingNode || srcMap.Kind != yaml.MappingNode { return } keep := make(map[string]struct{}, len(srcMap.Content)/2) for i := 0; i+1 < len(srcMap.Content); i += 2 { keyNode := srcMap.Content[i] if keyNode == nil { continue } key := strings.TrimSpace(keyNode.Value) if key == "" { continue } keep[key] = struct{}{} } for i := 0; i+1 < len(dstMap.Content); { keyNode := dstMap.Content[i] if keyNode == nil { i += 2 continue } key := strings.TrimSpace(keyNode.Value) if _, ok := keep[key]; !ok { dstMap.Content = append(dstMap.Content[:i], dstMap.Content[i+2:]...) continue } i += 2 } } // normalizeCollectionNodeStyles forces YAML collections to use block notation, keeping // lists and maps readable. Empty sequences retain flow style ([]) so empty list markers // remain compact. func normalizeCollectionNodeStyles(node *yaml.Node) { if node == nil { return } switch node.Kind { case yaml.MappingNode: node.Style = 0 for i := range node.Content { normalizeCollectionNodeStyles(node.Content[i]) } case yaml.SequenceNode: if len(node.Content) == 0 { node.Style = yaml.FlowStyle } else { node.Style = 0 } for i := range node.Content { normalizeCollectionNodeStyles(node.Content[i]) } default: // Scalars keep their existing style to preserve quoting } }