496 lines
14 KiB
Go
496 lines
14 KiB
Go
// Package util provides utility functions for the CLI Proxy API server.
|
|
// It includes helper functions for JSON manipulation, proxy configuration,
|
|
// and other common operations used across the application.
|
|
package util
|
|
|
|
import (
|
|
"bytes"
|
|
"crypto/sha256"
|
|
"encoding/hex"
|
|
"fmt"
|
|
"sort"
|
|
"strings"
|
|
|
|
log "github.com/sirupsen/logrus"
|
|
"github.com/tidwall/gjson"
|
|
"github.com/tidwall/sjson"
|
|
)
|
|
|
|
// Walk recursively traverses a JSON structure to find all occurrences of a specific field.
|
|
// It builds paths to each occurrence and adds them to the provided paths slice.
|
|
//
|
|
// Parameters:
|
|
// - value: The gjson.Result object to traverse
|
|
// - path: The current path in the JSON structure (empty string for root)
|
|
// - field: The field name to search for
|
|
// - paths: Pointer to a slice where found paths will be stored
|
|
//
|
|
// The function works recursively, building dot-notation paths to each occurrence
|
|
// of the specified field throughout the JSON structure.
|
|
func Walk(value gjson.Result, path, field string, paths *[]string) {
|
|
switch value.Type {
|
|
case gjson.JSON:
|
|
// For JSON objects and arrays, iterate through each child
|
|
value.ForEach(func(key, val gjson.Result) bool {
|
|
var childPath string
|
|
// Escape special characters for gjson/sjson path syntax
|
|
// . -> \.
|
|
// * -> \*
|
|
// ? -> \?
|
|
keyStr := key.String()
|
|
safeKey := escapeGJSONPathKey(keyStr)
|
|
|
|
if path == "" {
|
|
childPath = safeKey
|
|
} else {
|
|
childPath = path + "." + safeKey
|
|
}
|
|
if keyStr == field {
|
|
*paths = append(*paths, childPath)
|
|
}
|
|
Walk(val, childPath, field, paths)
|
|
return true
|
|
})
|
|
case gjson.String, gjson.Number, gjson.True, gjson.False, gjson.Null:
|
|
// Terminal types - no further traversal needed
|
|
}
|
|
}
|
|
|
|
// RenameKey renames a key in a JSON string by moving its value to a new key path
|
|
// and then deleting the old key path.
|
|
//
|
|
// Parameters:
|
|
// - jsonStr: The JSON string to modify
|
|
// - oldKeyPath: The dot-notation path to the key that should be renamed
|
|
// - newKeyPath: The dot-notation path where the value should be moved to
|
|
//
|
|
// Returns:
|
|
// - string: The modified JSON string with the key renamed
|
|
// - error: An error if the operation fails
|
|
//
|
|
// The function performs the rename in two steps:
|
|
// 1. Sets the value at the new key path
|
|
// 2. Deletes the old key path
|
|
func RenameKey(jsonStr, oldKeyPath, newKeyPath string) (string, error) {
|
|
value := gjson.Get(jsonStr, oldKeyPath)
|
|
|
|
if !value.Exists() {
|
|
return "", fmt.Errorf("old key '%s' does not exist", oldKeyPath)
|
|
}
|
|
|
|
interimJSON, errSet := sjson.SetRawBytes([]byte(jsonStr), newKeyPath, []byte(value.Raw))
|
|
if errSet != nil {
|
|
return "", fmt.Errorf("failed to set new key '%s': %w", newKeyPath, errSet)
|
|
}
|
|
|
|
finalJSON, errDelete := sjson.DeleteBytes(interimJSON, oldKeyPath)
|
|
if errDelete != nil {
|
|
return "", fmt.Errorf("failed to delete old key '%s': %w", oldKeyPath, errDelete)
|
|
}
|
|
|
|
return string(finalJSON), nil
|
|
}
|
|
|
|
// FixJSON converts non-standard JSON that uses single quotes for strings into
|
|
// RFC 8259-compliant JSON by converting those single-quoted strings to
|
|
// double-quoted strings with proper escaping.
|
|
//
|
|
// Examples:
|
|
//
|
|
// {'a': 1, 'b': '2'} => {"a": 1, "b": "2"}
|
|
// {"t": 'He said "hi"'} => {"t": "He said \"hi\""}
|
|
//
|
|
// Rules:
|
|
// - Existing double-quoted JSON strings are preserved as-is.
|
|
// - Single-quoted strings are converted to double-quoted strings.
|
|
// - Inside converted strings, any double quote is escaped (\").
|
|
// - Common backslash escapes (\n, \r, \t, \b, \f, \\) are preserved.
|
|
// - \' inside single-quoted strings becomes a literal ' in the output (no
|
|
// escaping needed inside double quotes).
|
|
// - Unicode escapes (\uXXXX) inside single-quoted strings are forwarded.
|
|
// - The function does not attempt to fix other non-JSON features beyond quotes.
|
|
func FixJSON(input string) string {
|
|
var out bytes.Buffer
|
|
|
|
inDouble := false
|
|
inSingle := false
|
|
escaped := false // applies within the current string state
|
|
|
|
// Helper to write a rune, escaping double quotes when inside a converted
|
|
// single-quoted string (which becomes a double-quoted string in output).
|
|
writeConverted := func(r rune) {
|
|
if r == '"' {
|
|
out.WriteByte('\\')
|
|
out.WriteByte('"')
|
|
return
|
|
}
|
|
out.WriteRune(r)
|
|
}
|
|
|
|
runes := []rune(input)
|
|
for i := 0; i < len(runes); i++ {
|
|
r := runes[i]
|
|
|
|
if inDouble {
|
|
out.WriteRune(r)
|
|
if escaped {
|
|
// end of escape sequence in a standard JSON string
|
|
escaped = false
|
|
continue
|
|
}
|
|
if r == '\\' {
|
|
escaped = true
|
|
continue
|
|
}
|
|
if r == '"' {
|
|
inDouble = false
|
|
}
|
|
continue
|
|
}
|
|
|
|
if inSingle {
|
|
if escaped {
|
|
// Handle common escape sequences after a backslash within a
|
|
// single-quoted string
|
|
escaped = false
|
|
switch r {
|
|
case 'n', 'r', 't', 'b', 'f', '/', '"':
|
|
// Keep the backslash and the character (except for '"' which
|
|
// rarely appears, but if it does, keep as \" to remain valid)
|
|
out.WriteByte('\\')
|
|
out.WriteRune(r)
|
|
case '\\':
|
|
out.WriteByte('\\')
|
|
out.WriteByte('\\')
|
|
case '\'':
|
|
// \' inside single-quoted becomes a literal '
|
|
out.WriteRune('\'')
|
|
case 'u':
|
|
// Forward \uXXXX if possible
|
|
out.WriteByte('\\')
|
|
out.WriteByte('u')
|
|
// Copy up to next 4 hex digits if present
|
|
for k := 0; k < 4 && i+1 < len(runes); k++ {
|
|
peek := runes[i+1]
|
|
// simple hex check
|
|
if (peek >= '0' && peek <= '9') || (peek >= 'a' && peek <= 'f') || (peek >= 'A' && peek <= 'F') {
|
|
out.WriteRune(peek)
|
|
i++
|
|
} else {
|
|
break
|
|
}
|
|
}
|
|
default:
|
|
// Unknown escape: preserve the backslash and the char
|
|
out.WriteByte('\\')
|
|
out.WriteRune(r)
|
|
}
|
|
continue
|
|
}
|
|
|
|
if r == '\\' { // start escape sequence
|
|
escaped = true
|
|
continue
|
|
}
|
|
if r == '\'' { // end of single-quoted string
|
|
out.WriteByte('"')
|
|
inSingle = false
|
|
continue
|
|
}
|
|
// regular char inside converted string; escape double quotes
|
|
writeConverted(r)
|
|
continue
|
|
}
|
|
|
|
// Outside any string
|
|
if r == '"' {
|
|
inDouble = true
|
|
out.WriteRune(r)
|
|
continue
|
|
}
|
|
if r == '\'' { // start of non-standard single-quoted string
|
|
inSingle = true
|
|
out.WriteByte('"')
|
|
continue
|
|
}
|
|
out.WriteRune(r)
|
|
}
|
|
|
|
// If input ended while still inside a single-quoted string, close it to
|
|
// produce the best-effort valid JSON.
|
|
if inSingle {
|
|
out.WriteByte('"')
|
|
}
|
|
|
|
return out.String()
|
|
}
|
|
|
|
func CanonicalToolName(name string) string {
|
|
canonical := strings.TrimSpace(name)
|
|
canonical = strings.TrimLeft(canonical, "_")
|
|
return strings.ToLower(canonical)
|
|
}
|
|
|
|
// ToolNameMapFromClaudeRequest returns a canonical-name -> original-name map extracted from a Claude request.
|
|
// It is used to restore exact tool name casing for clients that require strict tool name matching (e.g. Claude Code).
|
|
func ToolNameMapFromClaudeRequest(rawJSON []byte) map[string]string {
|
|
if len(rawJSON) == 0 || !gjson.ValidBytes(rawJSON) {
|
|
return nil
|
|
}
|
|
|
|
tools := gjson.GetBytes(rawJSON, "tools")
|
|
if !tools.Exists() || !tools.IsArray() {
|
|
return nil
|
|
}
|
|
|
|
toolResults := tools.Array()
|
|
out := make(map[string]string, len(toolResults))
|
|
tools.ForEach(func(_, tool gjson.Result) bool {
|
|
name := strings.TrimSpace(tool.Get("name").String())
|
|
if name == "" {
|
|
name = strings.TrimSpace(tool.Get("function.name").String())
|
|
}
|
|
if name == "" {
|
|
return true
|
|
}
|
|
key := CanonicalToolName(name)
|
|
if key == "" {
|
|
return true
|
|
}
|
|
if _, exists := out[key]; !exists {
|
|
out[key] = name
|
|
}
|
|
return true
|
|
})
|
|
|
|
if len(out) == 0 {
|
|
return nil
|
|
}
|
|
return out
|
|
}
|
|
|
|
func MapToolName(toolNameMap map[string]string, name string) string {
|
|
if name == "" || toolNameMap == nil {
|
|
return name
|
|
}
|
|
if mapped, ok := toolNameMap[CanonicalToolName(name)]; ok && mapped != "" {
|
|
return mapped
|
|
}
|
|
return name
|
|
}
|
|
|
|
// SanitizedFunctionNameMap builds an original-name → sanitized-name map from request tools.
|
|
// Exact duplicate names share a mapping. Distinct names that sanitize to the same value receive
|
|
// deterministic hash suffixes so every declaration remains addressable within the 64-byte limit.
|
|
func SanitizedFunctionNameMap(rawJSON []byte) map[string]string {
|
|
names := functionNamesFromRequest(rawJSON)
|
|
if len(names) == 0 {
|
|
return nil
|
|
}
|
|
|
|
uniqueNames := make(map[string]struct{}, len(names))
|
|
baseCounts := make(map[string]int, len(names))
|
|
for _, name := range names {
|
|
if name == "" {
|
|
continue
|
|
}
|
|
if _, exists := uniqueNames[name]; exists {
|
|
continue
|
|
}
|
|
uniqueNames[name] = struct{}{}
|
|
baseCounts[SanitizeFunctionName(name)]++
|
|
}
|
|
|
|
sortedNames := make([]string, 0, len(uniqueNames))
|
|
for name := range uniqueNames {
|
|
sortedNames = append(sortedNames, name)
|
|
}
|
|
sort.Strings(sortedNames)
|
|
|
|
out := make(map[string]string, len(sortedNames))
|
|
used := make(map[string]string, len(sortedNames))
|
|
for _, name := range sortedNames {
|
|
base := SanitizeFunctionName(name)
|
|
mapped := base
|
|
_, baseUsed := used[base]
|
|
if baseCounts[base] > 1 || baseUsed {
|
|
mapped = disambiguateSanitizedFunctionName(base, name, used)
|
|
}
|
|
out[name] = mapped
|
|
used[mapped] = name
|
|
}
|
|
if len(out) == 0 {
|
|
return nil
|
|
}
|
|
return out
|
|
}
|
|
|
|
// MapSanitizedFunctionName returns the request-specific sanitized name when available.
|
|
func MapSanitizedFunctionName(nameMap map[string]string, name string) string {
|
|
if mapped := nameMap[name]; mapped != "" {
|
|
return mapped
|
|
}
|
|
return SanitizeFunctionName(name)
|
|
}
|
|
|
|
// DisambiguatedToolNameMap builds a sanitized-name → original-name map using the
|
|
// same collision-aware mapping as SanitizedFunctionNameMap.
|
|
func DisambiguatedToolNameMap(rawJSON []byte) map[string]string {
|
|
forward := SanitizedFunctionNameMap(rawJSON)
|
|
if len(forward) == 0 {
|
|
return nil
|
|
}
|
|
|
|
out := make(map[string]string, len(forward))
|
|
for original, sanitized := range forward {
|
|
if sanitized != original {
|
|
out[sanitized] = original
|
|
}
|
|
}
|
|
if len(out) == 0 {
|
|
return nil
|
|
}
|
|
return out
|
|
}
|
|
|
|
// SanitizedToolNameMap builds the legacy sanitized-name → original-name map from
|
|
// top-level Claude-style tools. Collision-aware translators should use
|
|
// DisambiguatedToolNameMap instead.
|
|
func SanitizedToolNameMap(rawJSON []byte) map[string]string {
|
|
if len(rawJSON) == 0 || !gjson.ValidBytes(rawJSON) {
|
|
return nil
|
|
}
|
|
tools := gjson.GetBytes(rawJSON, "tools")
|
|
if !tools.IsArray() {
|
|
return nil
|
|
}
|
|
|
|
out := make(map[string]string)
|
|
tools.ForEach(func(_, tool gjson.Result) bool {
|
|
name := strings.TrimSpace(tool.Get("name").String())
|
|
if name == "" {
|
|
return true
|
|
}
|
|
sanitized := SanitizeFunctionName(name)
|
|
if sanitized == name {
|
|
return true
|
|
}
|
|
if existing, exists := out[sanitized]; !exists {
|
|
out[sanitized] = name
|
|
} else {
|
|
log.Warnf("sanitized tool name collision: %q and %q both map to %q, keeping first", existing, name, sanitized)
|
|
}
|
|
return true
|
|
})
|
|
if len(out) == 0 {
|
|
return nil
|
|
}
|
|
return out
|
|
}
|
|
|
|
func functionNamesFromRequest(rawJSON []byte) []string {
|
|
if len(rawJSON) == 0 || !gjson.ValidBytes(rawJSON) {
|
|
return nil
|
|
}
|
|
tools := gjson.GetBytes(rawJSON, "tools")
|
|
if !tools.IsArray() {
|
|
return nil
|
|
}
|
|
|
|
names := make([]string, 0, len(tools.Array()))
|
|
var collectTool func(gjson.Result)
|
|
collectDeclarations := func(declarations gjson.Result) {
|
|
if !declarations.IsArray() {
|
|
return
|
|
}
|
|
declarations.ForEach(func(_, declaration gjson.Result) bool {
|
|
if name := declaration.Get("name").String(); name != "" {
|
|
names = append(names, name)
|
|
}
|
|
return true
|
|
})
|
|
}
|
|
collectTool = func(tool gjson.Result) {
|
|
if nestedTools := tool.Get("tools"); nestedTools.IsArray() {
|
|
nestedTools.ForEach(func(_, nestedTool gjson.Result) bool {
|
|
collectTool(nestedTool)
|
|
return true
|
|
})
|
|
return
|
|
}
|
|
hasDeclarations := false
|
|
if declarations := tool.Get("functionDeclarations"); declarations.IsArray() {
|
|
collectDeclarations(declarations)
|
|
hasDeclarations = true
|
|
}
|
|
if declarations := tool.Get("function_declarations"); declarations.IsArray() {
|
|
collectDeclarations(declarations)
|
|
hasDeclarations = true
|
|
}
|
|
if hasDeclarations {
|
|
return
|
|
}
|
|
if name := tool.Get("function.name").String(); name != "" {
|
|
names = append(names, name)
|
|
return
|
|
}
|
|
if name := tool.Get("name").String(); name != "" {
|
|
names = append(names, name)
|
|
}
|
|
}
|
|
tools.ForEach(func(_, tool gjson.Result) bool {
|
|
collectTool(tool)
|
|
return true
|
|
})
|
|
return names
|
|
}
|
|
|
|
func disambiguateSanitizedFunctionName(base, original string, used map[string]string) string {
|
|
for attempt := 0; ; attempt++ {
|
|
digest := sha256.Sum256([]byte(fmt.Sprintf("%s\x00%d", original, attempt)))
|
|
suffix := "_" + hex.EncodeToString(digest[:6])
|
|
prefix := base
|
|
if maxPrefix := 64 - len(suffix); len(prefix) > maxPrefix {
|
|
prefix = prefix[:maxPrefix]
|
|
}
|
|
candidate := prefix + suffix
|
|
if _, exists := used[candidate]; !exists {
|
|
return candidate
|
|
}
|
|
}
|
|
}
|
|
|
|
// DeduplicateFunctionDeclarations removes duplicate named declarations while preserving order.
|
|
func DeduplicateFunctionDeclarations(raw []byte) []byte {
|
|
result := gjson.ParseBytes(raw)
|
|
if !result.IsArray() {
|
|
return raw
|
|
}
|
|
|
|
seen := make(map[string]struct{}, len(result.Array()))
|
|
parts := make([]string, 0, len(result.Array()))
|
|
for _, declaration := range result.Array() {
|
|
name := declaration.Get("name").String()
|
|
if name != "" {
|
|
if _, exists := seen[name]; exists {
|
|
continue
|
|
}
|
|
seen[name] = struct{}{}
|
|
}
|
|
parts = append(parts, declaration.Raw)
|
|
}
|
|
return []byte("[" + strings.Join(parts, ",") + "]")
|
|
}
|
|
|
|
// RestoreSanitizedToolName looks up a sanitized function name in the provided map
|
|
// and returns the original client-facing name. If no mapping exists, it returns
|
|
// the sanitized name unchanged.
|
|
func RestoreSanitizedToolName(toolNameMap map[string]string, sanitizedName string) string {
|
|
if sanitizedName == "" || toolNameMap == nil {
|
|
return sanitizedName
|
|
}
|
|
if original, ok := toolNameMap[sanitizedName]; ok {
|
|
return original
|
|
}
|
|
return sanitizedName
|
|
}
|