MuyueWorkspace/internal/agent/definitions.go

package agent

import (
	"context"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"sync"
	"time"
)

var ansiRegex = regexp.MustCompile(`\x1b\[[0-9;]*[a-zA-Z]|\x1b\][^\x07]*\x07|\x1b\][^\x1b]*\x1b\\|\x1b[()][AB012]|\[\]`)

func stripANSI(s string) string {
	return ansiRegex.ReplaceAllString(s, "")
}

var (
	sudoCache     bool
	sudoCacheSet  bool
	sudoCacheOnce sync.Once
)

func NeedsSudoPassword() bool {
	sudoCacheOnce.Do(func() {
		if os.Geteuid() == 0 {
			ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
			defer cancel()
			err := exec.CommandContext(ctx, "sudo", "-n", "true").Run()
			sudoCacheSet = true
			sudoCache = err != nil
		} else {
			sudoCache = true
			sudoCacheSet = true
		}
	})
	return sudoCache
}

type TerminalParams struct {
	Command string `json:"command" description:"The shell command to execute"`
	Timeout int    `json:"timeout,omitempty" description:"Timeout in seconds (default 60, max 300)"`
}

type TerminalResponse struct {
	Content     string `json:"content"`
	IsError     bool   `json:"is_error"`
	SudoBlocked bool   `json:"sudo_blocked,omitempty"`
	Command     string `json:"command,omitempty"`
}

func NewTerminalTool() (*ToolDefinition, error) {
	return NewTool("terminal",
		"Execute a shell command on the local system and return the output. Use for running builds, tests, git operations, package management, system info, or any CLI task. Commands run in the user's home directory by default. Long-running commands are auto-terminated.",
		func(ctx context.Context, p TerminalParams) (ToolResponse, error) {
			if p.Command == "" {
				return TextErrorResponse("command is required"), nil
			}

			if NeedsSudoPassword() {
				trimmed := strings.TrimSpace(p.Command)
				lower := strings.ToLower(trimmed)
				prefixBlocked := strings.HasPrefix(lower, "sudo ") || strings.HasPrefix(lower, "doas ") || strings.HasPrefix(lower, "run0 ") || strings.HasPrefix(lower, "pkexec ")
				anywhereBlocked := false
				blockedCmd := ""
				if !prefixBlocked {
					for _, kw := range []string{"sudo", "doas", "run0", "pkexec"} {
						for _, pattern := range []string{" " + kw + " ", "|" + kw + " ", ";" + kw + " ", "&&" + kw + " ", "||" + kw + " ", "`" + kw + " ", "$(" + kw + " "} {
							if strings.Contains(lower, pattern) {
								anywhereBlocked = true
								blockedCmd = kw
								break
							}
						}
						if anywhereBlocked {
							break
						}
					}
				}
				if prefixBlocked || anywhereBlocked {
					elevCmd := blockedCmd
					if prefixBlocked {
						elevCmd = strings.Fields(trimmed)[0]
					}
					return ToolResponse{
						Content: fmt.Sprintf("BLOCKED: Command '%s' requires elevated privileges (%s). Passwordless sudo is not available. Do NOT retry with sudo. Explain to the user that this command needs admin privileges and suggest an alternative, or tell them to run it manually in their terminal.", trimmed, elevCmd),
						IsError: true,
						Meta:    map[string]string{"sudo_blocked": "true", "command": trimmed},
					}, nil
				}
			}

			timeout := time.Duration(p.Timeout) * time.Second
			if timeout == 0 {
				timeout = 60 * time.Second
			}
			if timeout > 300*time.Second {
				timeout = 300 * time.Second
			}

			ctx, cancel := context.WithTimeout(ctx, timeout)
			defer cancel()

			shell := detectShell()

			cmd := exec.CommandContext(ctx, shell, "-c", p.Command)
			output, err := cmd.CombinedOutput()

			result := string(output)
			result = stripANSI(result)
			if len(result) > 10000 {
				result = result[:10000] + "\n... [truncated]"
			}

			if err != nil {
				return TextErrorResponse(fmt.Sprintf("Error: %v\n\n%s", err, result)), nil
			}

			return TextResponse(result), nil
		})
}

type CrushRunParams struct {
	Task       string `json:"task" description:"The task description for Crush to execute"`
	Timeout    int    `json:"timeout,omitempty" description:"Maximum execution time in seconds (default 1800, max 1800)"`
	Cwd        string `json:"cwd,omitempty" description:"Working directory in which to launch the agent (absolute path; falls back to user home)"`
	WSLDistro  string `json:"wsl_distro,omitempty" description:"On Windows host: WSL distribution to launch the agent in (e.g. 'Ubuntu')"`
	WSLUser    string `json:"wsl_user,omitempty" description:"On Windows host: WSL user to run the agent as"`
}

func NewCrushRunTool() (*ToolDefinition, error) {
	return NewTool("crush_run",
		"Delegate a complex coding task to the Crush AI agent. Crush has access to file editing, code search, bash execution, and other development tools. Use this for multi-step coding tasks like refactoring, debugging, implementing features, or code review. Optionally pass cwd to run in a specific directory, or wsl_distro/wsl_user to launch inside a WSL distribution under a specific user (Windows hosts only). Returns the agent's final output.",
		func(ctx context.Context, p CrushRunParams) (ToolResponse, error) {
			if p.Task == "" {
				return TextErrorResponse("task is required"), nil
			}

			timeout := time.Duration(p.Timeout) * time.Second
			if timeout == 0 {
				timeout = 1800 * time.Second
			}
			if timeout > 1800*time.Second {
				timeout = 1800 * time.Second
			}
			ctx, cancel := context.WithTimeout(ctx, timeout)
			defer cancel()

			cmd, prepErr := buildAgentCommand(ctx, "crush", []string{"run", p.Task}, p.Cwd, p.WSLDistro, p.WSLUser)
			if prepErr != nil {
				return TextErrorResponse(prepErr.Error()), nil
			}
			output, err := cmd.CombinedOutput()

			result := string(output)
			if len(result) > 15000 {
				result = result[:15000] + "\n... [truncated]"
			}

			if err != nil {
				errMsg := fmt.Sprintf("Crush error: %v", err)
				if ctx.Err() == context.DeadlineExceeded {
					errMsg = fmt.Sprintf("Crush timed out after %d seconds. Try splitting the task into smaller parts.", int(timeout.Seconds()))
				}
				if result != "" {
					errMsg += "\n\n" + result
				}
				return TextErrorResponse(errMsg), nil
			}

			return TextResponse(result), nil
		})
}

type ClaudeRunParams struct {
	Task       string `json:"task" description:"The task description for Claude Code to execute"`
	Timeout    int    `json:"timeout,omitempty" description:"Maximum execution time in seconds (default 1800, max 1800)"`
	Cwd        string `json:"cwd,omitempty" description:"Working directory in which to launch the agent (absolute path; falls back to user home)"`
	WSLDistro  string `json:"wsl_distro,omitempty" description:"On Windows host: WSL distribution to launch the agent in (e.g. 'Ubuntu')"`
	WSLUser    string `json:"wsl_user,omitempty" description:"On Windows host: WSL user to run the agent as"`
}

func NewClaudeRunTool() (*ToolDefinition, error) {
	return NewTool("claude_run",
		"Delegate a complex coding task to the Claude Code CLI agent. Claude has access to file editing, code search, bash execution. Use for multi-step coding tasks. Same cwd/wsl_distro/wsl_user options as crush_run.",
		func(ctx context.Context, p ClaudeRunParams) (ToolResponse, error) {
			if p.Task == "" {
				return TextErrorResponse("task is required"), nil
			}

			timeout := time.Duration(p.Timeout) * time.Second
			if timeout == 0 {
				timeout = 1800 * time.Second
			}
			if timeout > 1800*time.Second {
				timeout = 1800 * time.Second
			}
			ctx, cancel := context.WithTimeout(ctx, timeout)
			defer cancel()

			cmd, prepErr := buildAgentCommand(ctx, "claude", []string{"-p", p.Task}, p.Cwd, p.WSLDistro, p.WSLUser)
			if prepErr != nil {
				return TextErrorResponse(prepErr.Error()), nil
			}
			output, err := cmd.CombinedOutput()

			result := string(output)
			if len(result) > 15000 {
				result = result[:15000] + "\n... [truncated]"
			}

			if err != nil {
				errMsg := fmt.Sprintf("Claude error: %v", err)
				if ctx.Err() == context.DeadlineExceeded {
					errMsg = fmt.Sprintf("Claude timed out after %d seconds. Try splitting the task into smaller parts.", int(timeout.Seconds()))
				}
				if result != "" {
					errMsg += "\n\n" + result
				}
				return TextErrorResponse(errMsg), nil
			}

			return TextResponse(result), nil
		})
}

type ReadFileParams struct {
	Path  string `json:"path" description:"Absolute or relative path to the file to read"`
	Offset int   `json:"offset,omitempty" description:"Line number to start reading from (0-based, default 0)"`
	Limit  int   `json:"limit,omitempty" description:"Maximum number of lines to read (default 200, max 2000)"`
}

func NewReadFileTool() (*ToolDefinition, error) {
	return NewTool("read_file",
		"Read file contents from the local filesystem. Returns the file content with line numbers. Supports offset/limit for reading specific sections of large files.",
		func(ctx context.Context, p ReadFileParams) (ToolResponse, error) {
			if p.Path == "" {
				return TextErrorResponse("path is required"), nil
			}

			expanded := expandHome(p.Path)
			data, err := readFileLimited(expanded, p.Offset, p.Limit)
			if err != nil {
				return TextErrorResponse(fmt.Sprintf("read error: %v", err)), nil
			}

			return TextResponse(data), nil
		})
}

type ListFilesParams struct {
	Path  string `json:"path,omitempty" description:"Directory path to list (default: user home)"`
	Depth int    `json:"depth,omitempty" description:"Maximum depth to traverse (default 1, max 3)"`
}

func NewListFilesTool() (*ToolDefinition, error) {
	return NewTool("list_files",
		"List files and directories at a given path. Shows directory tree structure with file names. Useful for exploring project structure or finding specific files.",
		func(ctx context.Context, p ListFilesParams) (ToolResponse, error) {
			dir := expandHome(p.Path)
			if dir == "" {
				dir, _ = osUserHomeDir()
			}

			if p.Depth <= 0 {
				p.Depth = 1
			}
			if p.Depth > 3 {
				p.Depth = 3
			}

			result, err := listDirTree(dir, p.Depth, 0)
			if err != nil {
				return TextErrorResponse(fmt.Sprintf("list error: %v", err)), nil
			}

			return TextResponse(result), nil
		})
}

type SearchFilesParams struct {
	Pattern string `json:"pattern" description:"Search pattern (supports * and ? glob wildcards)"`
	Path    string `json:"path,omitempty" description:"Directory to search in (default: current directory)"`
}

func NewSearchFilesTool() (*ToolDefinition, error) {
	return NewTool("search_files",
		"Search for files by name pattern using glob syntax. Use * for any characters, ** for recursive matching. Returns matching file paths sorted by name.",
		func(ctx context.Context, p SearchFilesParams) (ToolResponse, error) {
			if p.Pattern == "" {
				return TextErrorResponse("pattern is required"), nil
			}

			dir := expandHome(p.Path)
			if dir == "" {
				dir = "."
			}

			matches, err := filepath.Glob(filepath.Join(dir, p.Pattern))
			if err != nil {
				return TextErrorResponse(fmt.Sprintf("glob error: %v", err)), nil
			}

			if len(matches) == 0 {
				return TextResponse("No files found matching pattern."), nil
			}

			if len(matches) > 100 {
				matches = matches[:100]
			}

			var result strings.Builder
			for _, m := range matches {
				result.WriteString(m)
				result.WriteString("\n")
			}

			return TextResponse(result.String()), nil
		})
}

type GrepContentParams struct {
	Pattern string `json:"pattern" description:"Text pattern to search for in file contents"`
	Path    string `json:"path,omitempty" description:"Directory to search in (default: current directory)"`
	Include string `json:"include,omitempty" description:"File extension filter, e.g. '*.go' or '*.{js,ts}'"`
}

func NewGrepContentTool() (*ToolDefinition, error) {
	return NewTool("grep_content",
		"Search for text patterns inside file contents. Returns matching lines with file paths and line numbers. Use include to filter by file extension.",
		func(ctx context.Context, p GrepContentParams) (ToolResponse, error) {
			if p.Pattern == "" {
				return TextErrorResponse("pattern is required"), nil
			}

			dir := expandHome(p.Path)
			if dir == "" {
				dir = "."
			}

			result, err := grepFiles(dir, p.Pattern, p.Include)
			if err != nil {
				return TextErrorResponse(fmt.Sprintf("grep error: %v", err)), nil
			}

			if result == "" {
				return TextResponse("No matches found."), nil
			}

			return TextResponse(result), nil
		})
}

type GetConfigParams struct {
	Section string `json:"section,omitempty" description:"Config section to retrieve: 'providers', 'profile', 'tools', 'terminal', 'all' (default: 'all')"`
}

func NewGetConfigTool() (*ToolDefinition, error) {
	return NewTool("get_config",
		"Read the Muyue configuration. Returns provider settings, profile info, installed tools, terminal config, etc. Use section parameter to get a specific part, or 'all' for the full config.",
		func(ctx context.Context, p GetConfigParams) (ToolResponse, error) {
			return getConfigSection(p.Section), nil
		})
}

type SetProviderParams struct {
	Name    string `json:"name" description:"Provider name (e.g. 'openai', 'anthropic', 'ollama')"`
	APIKey  string `json:"api_key,omitempty" description:"API key for the provider"`
	BaseURL string `json:"base_url,omitempty" description:"Custom base URL for the provider API"`
	Model   string `json:"model,omitempty" description:"Model identifier to use"`
	Active  *bool  `json:"active,omitempty" description:"Set to true to make this the active provider"`
}

func NewSetProviderTool() (*ToolDefinition, error) {
	return NewTool("set_provider",
		"Configure an AI provider in Muyue settings. Can create, update, or activate a provider. API keys are automatically encrypted. Set active=true to switch to this provider.",
		func(ctx context.Context, p SetProviderParams) (ToolResponse, error) {
			if p.Name == "" {
				return TextErrorResponse("name is required"), nil
			}

			return setProviderConfig(p), nil
		})
}

type ManageSSHParams struct {
	Action  string `json:"action" description:"Action to perform: 'list', 'add', 'remove'"`
	Name    string `json:"name,omitempty" description:"Connection name (required for add/remove)"`
	Host    string `json:"host,omitempty" description:"SSH host (required for add)"`
	Port    int    `json:"port,omitempty" description:"SSH port (default: 22)"`
	User    string `json:"user,omitempty" description:"SSH username (required for add)"`
	KeyPath string `json:"key_path,omitempty" description:"Path to SSH private key"`
}

func NewManageSSHTool() (*ToolDefinition, error) {
	return NewTool("manage_ssh",
		"Manage SSH connections configured in Muyue. List existing connections, add new ones, or remove connections. SSH configs are persisted to the Muyue config file.",
		func(ctx context.Context, p ManageSSHParams) (ToolResponse, error) {
			if p.Action == "" {
				return TextErrorResponse("action is required (list, add, remove)"), nil
			}

			return manageSSHAction(p), nil
		})
}

type WebFetchParams struct {
	URL string `json:"url" description:"The URL to fetch content from"`
}

func NewWebFetchTool() (*ToolDefinition, error) {
	return NewTool("web_fetch",
		"Fetch content from a URL and return the text. Useful for reading documentation, APIs, or web resources. Only HTTP/HTTPS URLs are supported.",
		func(ctx context.Context, p WebFetchParams) (ToolResponse, error) {
			if p.URL == "" {
				return TextErrorResponse("url is required"), nil
			}

			return fetchURL(p.URL), nil
		})
}

func DefaultRegistry() *Registry {
	r := NewRegistry()

	tools := []*ToolDefinition{
		must(NewTerminalTool()),
		must(NewCrushRunTool()),
		must(NewClaudeRunTool()),
		must(NewReadFileTool()),
		must(NewListFilesTool()),
		must(NewSearchFilesTool()),
		must(NewGrepContentTool()),
		must(NewGetConfigTool()),
		must(NewSetProviderTool()),
		must(NewManageSSHTool()),
		must(NewWebFetchTool()),
	}

	for _, t := range tools {
		if err := r.Register(t); err != nil {
			panic(err)
		}
	}

	return r
}

func must(t *ToolDefinition, err error) *ToolDefinition {
	if err != nil {
		panic(err)
	}
	return t
}