ZovOS_AI/internal/tools/terminal/terminal.go

package terminal

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"os/exec"
	"strings"
	"time"

	"git.bebrik.xyz/Dan4ick/ZovOS_AI/internal/ai"
	"git.bebrik.xyz/Dan4ick/ZovOS_AI/internal/config"
	"git.bebrik.xyz/Dan4ick/ZovOS_AI/internal/tools"
)

type runCommandArgs struct {
	Command string `json:"command"`
	Cwd     string `json:"cwd,omitempty"`
}

// RegisterTools registers terminal-related tools in the registry.
func RegisterTools(registry *tools.Registry, cfg config.TerminalToolConfig, log *slog.Logger) {
	if !cfg.Enabled {
		return
	}

	registry.Register(tools.ToolDef{
		Tool: ai.Tool{
			Name:        "run_command",
			Description: "Execute a shell command in the terminal and return its output. Use this to run any command-line operations.",
			Parameters: map[string]interface{}{
				"type": "object",
				"properties": map[string]interface{}{
					"command": map[string]interface{}{
						"type":        "string",
						"description": "The shell command to execute",
					},
					"cwd": map[string]interface{}{
						"type":        "string",
						"description": "Working directory for the command (optional)",
					},
				},
				"required": []string{"command"},
			},
		},
		Handler: func(ctx context.Context, args json.RawMessage) (string, error) {
			var a runCommandArgs
			if err := json.Unmarshal(args, &a); err != nil {
				return "", fmt.Errorf("parsing args: %w", err)
			}
			return runCommand(ctx, a, cfg, log)
		},
	})
}

func runCommand(ctx context.Context, args runCommandArgs, cfg config.TerminalToolConfig, log *slog.Logger) (string, error) {
	// Check for blocked commands
	cmdLower := strings.ToLower(args.Command)
	for _, blocked := range cfg.BlockedCommands {
		if strings.Contains(cmdLower, strings.ToLower(blocked)) {
			return "", fmt.Errorf("command blocked by security policy: %s", blocked)
		}
	}

	timeout := time.Duration(cfg.TimeoutSeconds) * time.Second
	ctx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	shell := cfg.Shell
	if shell == "" {
		shell = "/bin/bash"
	}

	cmd := exec.CommandContext(ctx, shell, "-c", args.Command)
	if args.Cwd != "" {
		cmd.Dir = args.Cwd
	}

	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr

	log.Info("running command", "command", args.Command, "cwd", args.Cwd)

	err := cmd.Run()

	var result strings.Builder
	if stdout.Len() > 0 {
		result.WriteString("STDOUT:\n")
		result.WriteString(stdout.String())
	}
	if stderr.Len() > 0 {
		if result.Len() > 0 {
			result.WriteString("\n")
		}
		result.WriteString("STDERR:\n")
		result.WriteString(stderr.String())
	}

	if err != nil {
		if result.Len() > 0 {
			result.WriteString("\n")
		}
		result.WriteString(fmt.Sprintf("EXIT ERROR: %v", err))
	}

	if result.Len() == 0 {
		result.WriteString("(no output)")
	}

	// Truncate very long output
	output := result.String()
	const maxLen = 10000
	if len(output) > maxLen {
		output = output[:maxLen] + "\n... (output truncated)"
	}

	return output, nil
}