iop/apps/edge/internal/openai/tool_command_normalization.go

package openai

import (
	"strings"
)

type commandArgSchema struct {
	allowsString bool
	allowsObject bool
	objectProps  map[string]any
}

func normalizeCommandArgument(command any, schema commandArgSchema) (any, bool) {
	if text, ok := command.(string); ok {
		text = strings.TrimSpace(text)
		if text == "" {
			return command, false
		}
		if schema.allowsObject && !schema.allowsString && !commandStringNeedsShell(text) {
			if structured, ok := structuredCommandFromShellWords(text); ok {
				return structured, true
			}
		}
		return text, text != command
	}
	m, ok := command.(map[string]any)
	if !ok {
		return command, false
	}
	commandText, ok := m["command"].(string)
	if !ok {
		return command, false
	}
	if args, ok := stringSliceArgument(m["args"]); ok {
		commandText = strings.TrimSpace(commandText)
		if schema.allowsString {
			return renderShellCommand(commandText, args), true
		}
		if commandExecutableNeedsShell(commandText) || commandArgsNeedShell(args) {
			if schema.allowsString || !schema.allowsObject {
				return renderShellCommand(commandText, args), true
			}
			return cleanedStructuredCommand(commandText, args, schema.objectProps), true
		}
		if schema.allowsObject || !schema.allowsString {
			cleaned := cleanedStructuredCommand(commandText, args, schema.objectProps)
			return cleaned, !commandStructuredArgsAlreadyClean(m, args) || commandObjectHasStrippedKeys(m, schema.objectProps)
		}
		return renderShellCommand(commandText, args), true
	}
	commandText = strings.TrimSpace(commandText)
	if commandText != "" {
		if schema.allowsString {
			return commandText, true
		}
		if commandStringNeedsShell(commandText) {
			return commandText, true
		}
		if (schema.allowsObject || !schema.allowsString) && !commandStringNeedsShell(commandText) {
			structured, ok := structuredCommandFromShellWords(commandText)
			if ok {
				if len(schema.objectProps) > 0 {
					structured = filterMapToSchemaProperties(structured, schema.objectProps)
				}
				return structured, true
			}
		}
		if structured, ok := structuredCommandFromShellWords(commandText); ok && schema.allowsObject && !schema.allowsString {
			return structured, true
		}
		return commandText, true
	}
	return command, false
}

func shouldStripToolArgumentKey(key string, props map[string]any) bool {
	if len(props) > 0 {
		return !schemaHasProperty(props, key)
	}
	return key == "description" || key == "runInTerminal"
}

func schemaHasProperty(props map[string]any, key string) bool {
	_, ok := props[key]
	return ok
}

func schemaSingleCommandField(props map[string]any) string {
	for _, key := range []string{"command", "cmd"} {
		if schemaAllowsType(props[key], "string") {
			return key
		}
	}
	return ""
}

func firstCommandAsShellString(commands []any) (string, bool) {
	if len(commands) == 0 {
		return "", false
	}
	return commandAsShellString(commands[0])
}

func commandAsShellString(command any) (string, bool) {
	switch v := command.(type) {
	case string:
		text := strings.TrimSpace(v)
		return text, text != ""
	case map[string]any:
		commandText, ok := v["command"].(string)
		if !ok {
			return "", false
		}
		args, _ := stringSliceArgument(v["args"])
		text := renderShellCommand(commandText, args)
		return text, strings.TrimSpace(text) != ""
	default:
		return "", false
	}
}

func cleanedStructuredCommand(commandText string, args []string, props map[string]any) map[string]any {
	cleaned := map[string]any{
		"command": commandText,
	}
	if len(args) > 0 {
		cleaned["args"] = args
	}
	if len(props) > 0 {
		return filterMapToSchemaProperties(cleaned, props)
	}
	return cleaned
}

func filterMapToSchemaProperties(m map[string]any, props map[string]any) map[string]any {
	out := make(map[string]any, len(m))
	for key, value := range m {
		if schemaHasProperty(props, key) {
			out[key] = value
		}
	}
	return out
}

func commandObjectHasStrippedKeys(m map[string]any, props map[string]any) bool {
	for key := range m {
		switch key {
		case "command", "args":
			continue
		default:
			if len(props) == 0 || !schemaHasProperty(props, key) {
				return true
			}
		}
	}
	return false
}

func stringSliceArgument(value any) ([]string, bool) {
	if value == nil {
		return nil, false
	}
	raw, ok := value.([]any)
	if !ok {
		return nil, false
	}
	out := make([]string, 0, len(raw))
	for _, item := range raw {
		s, ok := item.(string)
		if !ok {
			return nil, false
		}
		out = append(out, s)
	}
	return out, true
}

func commandStructuredArgsAlreadyClean(m map[string]any, args []string) bool {
	if len(m) != 2 {
		return false
	}
	rawArgs, ok := m["args"].([]any)
	if !ok || len(rawArgs) != len(args) {
		return false
	}
	for i, arg := range args {
		if rawArgs[i] != arg {
			return false
		}
	}
	return true
}

func shellJoinArgs(args []string) string {
	if len(args) == 0 {
		return ""
	}
	quoted := make([]string, 0, len(args))
	for _, arg := range args {
		quoted = append(quoted, shellQuoteArg(arg))
	}
	return " " + strings.Join(quoted, " ")
}

func renderShellCommand(command string, args []string) string {
	return strings.TrimSpace(command) + shellJoinArgs(args)
}

func shellQuoteArg(arg string) string {
	if arg == "" {
		return "''"
	}
	if isShellOperatorToken(arg) {
		return arg
	}
	if !strings.ContainsAny(arg, " \t\n\r|&;<>(){}[]*$`'\"\\") {
		return arg
	}
	return "'" + strings.ReplaceAll(arg, "'", "'\\''") + "'"
}

func structuredCommandFromShellWords(text string) (map[string]any, bool) {
	words := strings.Fields(text)
	if len(words) == 0 {
		return nil, false
	}
	out := map[string]any{
		"command": words[0],
	}
	if len(words) > 1 {
		out["args"] = words[1:]
	}
	return out, true
}

func commandExecutableNeedsShell(text string) bool {
	text = strings.TrimSpace(text)
	return text == "" || strings.ContainsAny(text, " \t\n\r|&;<>(){}[]*$`'\"\\") || commandStringNeedsShell(text)
}

func commandArgsNeedShell(args []string) bool {
	for _, arg := range args {
		if isShellOperatorToken(arg) {
			return true
		}
	}
	return false
}

func isShellOperatorToken(arg string) bool {
	switch arg {
	case "&&", "||", "|", "|&", ";", "&", ">", ">>", "<", "<<", "<<<", "2>", "2>>", "&>", "2>&1":
		return true
	default:
		return false
	}
}

func commandStringNeedsShell(text string) bool {
	text = strings.TrimSpace(text)
	if text == "" {
		return false
	}
	if strings.ContainsAny(text, "\n\r|&;<>(){}[]*$`'\"\\") {
		return true
	}
	first := firstShellWord(text)
	switch first {
	case "alias", "bg", "cd", "command", "export", "fg", "hash", "jobs", "read", "set", "source", "test", "type", "ulimit", "umask", "unalias", "unset", "[", "[[", ":":
		return true
	default:
		return false
	}
}

func firstShellWord(s string) string {
	for i, r := range s {
		switch r {
		case ' ', '\t', '\n', '\r':
			return s[:i]
		}
	}
	return s
}