MuyueWorkspace/internal/agent/definitions.go

package agent

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

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)
				if strings.HasPrefix(lower, "sudo ") || strings.HasPrefix(lower, "doas ") || strings.HasPrefix(lower, "run0 ") || strings.HasPrefix(lower, "pkexec ") {
					return ToolResponse{
						Content: fmt.Sprintf("BLOCKED: Command '%s' requires elevated privileges (%s). The current user is not root. 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, strings.Fields(trimmed)[0]),
						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)
			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"`
}

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. Returns the agent's final output.",
		func(ctx context.Context, p CrushRunParams) (ToolResponse, error) {
			if p.Task == "" {
				return TextErrorResponse("task is required"), nil
			}

			ctx, cancel := context.WithTimeout(ctx, 300*time.Second)
			defer cancel()

			cmd := exec.CommandContext(ctx, "crush", "run", p.Task)
			output, err := cmd.CombinedOutput()

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

			if err != nil {
				return TextErrorResponse(fmt.Sprintf("Crush error: %v\n\n%s", err, result)), 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(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
}