Agent-Native/runtime/tools/exec.go

package tools

import (
	"context"
	"encoding/json"
	"fmt"
	"math/rand"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"
	"time"
	"unicode/utf8"
)

// ExecConfig 声明执行类工具的运行条件，由组合根显式提供；
// 未提供（nil）时 run_command / run_python 一律不注册（fail-closed）。
// 注意：工作区路径约束不等于 OS 沙箱——进程仍以当前用户权限运行，
// 隔离需由宿主环境提供（容器/沙箱），本文件只做时限、输出、env 与审计收敛。
type ExecConfig struct {
	// Bash / Python 是解释器绝对路径；留空则对应工具不注册。
	Bash   string
	Python string
	// Env 传给子进程的完整白名单环境（KEY=VALUE）；不含本进程任何隐式变量。
	Env []string
	// ScriptDir 落盘脚本目录（相对 root，默认 .agent/scripts）；空用默认。
	ScriptDir string
}

const (
	execDefaultTimeout = 60 * time.Second
	execMaxTimeout     = 300 * time.Second
	execMaxOutput      = 64 << 10
)

// blockedCommandMarkers 最小命令黑名单：明显提权/自毁模式直接拒绝。
// 这不是沙箱，只是防呆；真实隔离靠宿主环境。
var blockedCommandMarkers = []string{
	"sudo ", "sudo\t", "sudo\n",
	"rm -rf /", "rm -fr /",
	":(){:|:&};:",
	"mkfs", "shutdown", "reboot",
}

// RegisterExecTools 注册 run_command / run_python（解释器未配置的项跳过注册）。
func RegisterExecTools(r *Registry, root string, cfg ExecConfig) {
	if strings.TrimSpace(cfg.Bash) != "" {
		r.Register(&runCommandTool{root: root, cfg: cfg})
	}
	if strings.TrimSpace(cfg.Python) != "" {
		r.Register(&runPythonTool{root: root, cfg: cfg})
	}
}

type runCommandTool struct {
	root string
	cfg  ExecConfig
}

func (t *runCommandTool) Name() string { return "run_command" }
func (t *runCommandTool) Description() string {
	return "Run a bash command in the workspace root (process group, default 60s timeout, 64 KiB output cap). Non-zero exit and timeout are reported as text. NOT an OS sandbox."
}
func (t *runCommandTool) Parameters() json.RawMessage {
	return json.RawMessage(`{"type":"object","required":["command"],"properties":{"command":{"type":"string"},"timeout_seconds":{"type":"integer","description":"1..300, default 60"}}}`)
}

func (t *runCommandTool) Execute(ctx context.Context, args json.RawMessage) (Result, error) {
	var obj struct {
		Command        string `json:"command"`
		TimeoutSeconds int    `json:"timeout_seconds"`
	}
	if err := json.Unmarshal(args, &obj); err != nil {
		return Result{}, err
	}
	if strings.TrimSpace(obj.Command) == "" {
		return Result{Content: "command is empty", IsError: true}, nil
	}
	if blocked(obj.Command) {
		return Result{Content: "command rejected by blocklist (sudo / destructive patterns are not allowed)", IsError: true}, nil
	}
	timeout := clampTimeout(obj.TimeoutSeconds)
	out, timedOut, err := runProcess(ctx, t.cfg.Bash, []string{"-c", obj.Command}, t.root, t.cfg.Env, timeout)
	content := formatProcessOutput(out, timedOut, err)
	return Result{Content: content, IsError: timedOut || err != nil}, nil
}

type runPythonTool struct {
	root string
	cfg  ExecConfig
}

func (t *runPythonTool) Name() string { return "run_python" }
func (t *runPythonTool) Description() string {
	return "Run a Python script: the code is persisted under .agent/scripts/ in the workspace and executed with the configured python3 interpreter (default 60s timeout, 64 KiB output cap)."
}
func (t *runPythonTool) Parameters() json.RawMessage {
	return json.RawMessage(`{"type":"object","required":["code"],"properties":{"code":{"type":"string"},"timeout_seconds":{"type":"integer","description":"1..300, default 60"}}}`)
}

func (t *runPythonTool) Execute(ctx context.Context, args json.RawMessage) (Result, error) {
	var obj struct {
		Code           string `json:"code"`
		TimeoutSeconds int    `json:"timeout_seconds"`
	}
	if err := json.Unmarshal(args, &obj); err != nil {
		return Result{}, err
	}
	if strings.TrimSpace(obj.Code) == "" {
		return Result{Content: "code is empty", IsError: true}, nil
	}
	dir := t.cfg.ScriptDir
	if dir == "" {
		dir = ".agent/scripts"
	}
	scriptDir := filepath.Join(t.root, filepath.FromSlash(dir))
	if err := os.MkdirAll(scriptDir, 0o755); err != nil {
		return Result{Content: err.Error(), IsError: true}, nil
	}
	name := fmt.Sprintf("%s-%06d.py", time.Now().UTC().Format("20060102T150405"), rand.Intn(1000000))
	script := filepath.Join(scriptDir, name)
	if err := os.WriteFile(script, []byte(obj.Code), 0o644); err != nil {
		return Result{Content: err.Error(), IsError: true}, nil
	}
	timeout := clampTimeout(obj.TimeoutSeconds)
	out, timedOut, err := runProcess(ctx, t.cfg.Python, []string{script}, t.root, t.cfg.Env, timeout)
	content := formatProcessOutput(out, timedOut, err)
	content = "script: " + filepath.ToSlash(filepath.Join(dir, name)) + "\n" + content
	return Result{Content: content, IsError: timedOut || err != nil}, nil
}

// runProcess 启动独立进程组的子进程；ctx 取消或超时先到者杀整组（SIGKILL）。
func runProcess(ctx context.Context, interpreter string, argv []string, dir string, env []string, timeout time.Duration) (out []byte, timedOut bool, err error) {
	runCtx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()
	cmd := exec.CommandContext(runCtx, interpreter, argv...)
	cmd.Dir = dir
	cmd.Env = env
	// 独立进程组：超时/取消时杀 -PGID，确保孙进程一并退出。
	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
	cmd.Cancel = func() error {
		if cmd.Process != nil {
			return syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
		}
		return nil
	}
	cmd.WaitDelay = 3 * time.Second
	out, err = cmd.CombinedOutput()
	timedOut = runCtx.Err() != nil && (ctx.Err() == nil || runCtx.Err() != ctx.Err())
	return out, timedOut, err
}

func clampTimeout(seconds int) time.Duration {
	if seconds <= 0 {
		return execDefaultTimeout
	}
	if seconds > int(execMaxTimeout.Seconds()) {
		return execMaxTimeout
	}
	return time.Duration(seconds) * time.Second
}

func blocked(command string) bool {
	lower := strings.ToLower(command)
	for _, m := range blockedCommandMarkers {
		if strings.Contains(lower, m) {
			return true
		}
	}
	return false
}

// formatProcessOutput 截断为头尾各半（总上限 execMaxOutput）并标注退出状态。
func formatProcessOutput(out []byte, timedOut bool, err error) string {
	text := string(out)
	if len(out) > execMaxOutput {
		half := execMaxOutput / 2
		dropped := len(out) - execMaxOutput
		head := CutUTF8(string(out), half)
		tail := string(out[len(out)-half:])
		for len(tail) > 0 && !utf8.RuneStart(tail[0]) {
			tail = tail[1:]
		}
		text = head + fmt.Sprintf("\n[... truncated %d bytes ...]\n", dropped) + tail
	}
	switch {
	case timedOut:
		text += "\n[timeout: process group killed]"
	case err != nil:
		if ee, ok := err.(*exec.ExitError); ok {
			text += fmt.Sprintf("\n[exit code %d]", ee.ExitCode())
		} else {
			text += "\n[error] " + err.Error()
		}
	default:
		text += "\n[exit code 0]"
	}
	return text
}