Agent-Native/runtime/tools/exec.go

package tools

import (
	"context"
	"encoding/json"
	"fmt"
	"math/rand"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"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）；不含本进程任何隐式变量。
	// nil 在注册/启动时归一为空白名单（不继承宿主环境）；显式空切片同样为空白名单。
	Env []string
	// ScriptDir 落盘脚本目录（相对 root，默认 .agent/scripts）；空用默认。
	ScriptDir string
}

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

// normalizeExecEnv maps a nil env to an empty whitelist baseline so os/exec
// does not inherit the host environment. An explicit empty or non-nil slice
// is returned unchanged.
func normalizeExecEnv(env []string) []string {
	if env == nil {
		return []string{}
	}
	return env
}

// 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) {
	cfg.Env = normalizeExecEnv(cfg.Env)
	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, total, timedOut, err := runProcess(ctx, t.cfg.Bash, []string{"-c", obj.Command}, t.root, t.cfg.Env, timeout)
	content := formatProcessOutput(out, total, 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, total, timedOut, err := runProcess(ctx, t.cfg.Python, []string{script}, t.root, t.cfg.Env, timeout)
	content := formatProcessOutput(out, total, timedOut, err)
	content = "script: " + filepath.ToSlash(filepath.Join(dir, name)) + "\n" + content
	return Result{Content: content, IsError: timedOut || err != nil}, nil
}

// runProcess 启动独立进程组的子进程；ctx 取消或超时先到者杀整组（SIGKILL）。
// 合并输出持续排入有界头尾缓冲，避免 CombinedOutput 无界增长；Wait 回收子进程。
func runProcess(ctx context.Context, interpreter string, argv []string, dir string, env []string, timeout time.Duration) (out []byte, total int, timedOut bool, err error) {
	runCtx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()
	cmd := exec.CommandContext(runCtx, interpreter, argv...)
	cmd.Dir = dir
	cmd.Env = normalizeExecEnv(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
	buf := newHeadTailBuffer(execMaxOutput)
	cmd.Stdout = buf
	cmd.Stderr = buf
	err = cmd.Run()
	out, total = buf.Snapshot()
	timedOut = runCtx.Err() != nil && (ctx.Err() == nil || runCtx.Err() != ctx.Err())
	return out, total, 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
}

// headTailBuffer keeps a fixed-size head and a ring-buffer tail while counting
// all bytes written, so chatty children cannot grow memory without bound.
type headTailBuffer struct {
	mu         sync.Mutex
	max        int
	half       int
	head       []byte
	tail       []byte
	tailPos    int
	tailFilled int
	total      int
}

func newHeadTailBuffer(max int) *headTailBuffer {
	if max < 2 {
		max = 2
	}
	half := max / 2
	return &headTailBuffer{
		max:  max,
		half: half,
		head: make([]byte, 0, half),
		tail: make([]byte, half),
	}
}

func (b *headTailBuffer) Write(p []byte) (int, error) {
	b.mu.Lock()
	defer b.mu.Unlock()
	n := len(p)
	for len(p) > 0 {
		if len(b.head) < b.half {
			take := b.half - len(b.head)
			if take > len(p) {
				take = len(p)
			}
			b.head = append(b.head, p[:take]...)
			b.total += take
			p = p[take:]
			continue
		}
		// Head is full: remaining bytes rotate through the fixed tail ring.
		for i := 0; i < len(p); i++ {
			b.tail[b.tailPos] = p[i]
			b.tailPos = (b.tailPos + 1) % b.half
			if b.tailFilled < b.half {
				b.tailFilled++
			}
		}
		b.total += len(p)
		break
	}
	return n, nil
}

// Snapshot returns retained bytes and the total byte count observed.
// When total <= max the slice is the full stream; when truncated it is
// head||tail of length max (callers use total for the dropped-bytes marker).
func (b *headTailBuffer) Snapshot() ([]byte, int) {
	b.mu.Lock()
	defer b.mu.Unlock()
	if b.total <= b.half {
		out := make([]byte, len(b.head))
		copy(out, b.head)
		return out, b.total
	}
	tail := b.tailBytesLocked()
	out := make([]byte, 0, len(b.head)+len(tail))
	out = append(out, b.head...)
	out = append(out, tail...)
	return out, b.total
}

func (b *headTailBuffer) tailBytesLocked() []byte {
	if b.tailFilled == 0 {
		return nil
	}
	if b.tailFilled < b.half {
		out := make([]byte, b.tailFilled)
		copy(out, b.tail[:b.tailFilled])
		return out
	}
	out := make([]byte, b.half)
	n := copy(out, b.tail[b.tailPos:])
	copy(out[n:], b.tail[:b.tailPos])
	return out
}

// formatProcessOutput 截断为头尾各半（总上限 execMaxOutput）并标注退出状态。
// out 为有界捕获结果；total 为子进程实际写出字节数（可大于 len(out)）。
func formatProcessOutput(out []byte, total int, timedOut bool, err error) string {
	text := string(out)
	if total > execMaxOutput {
		half := execMaxOutput / 2
		dropped := total - execMaxOutput
		head := CutUTF8(string(out), half)
		tail := ""
		if len(out) > half {
			tail = string(out[len(out)-half:])
		} else if len(out) > len(head) {
			tail = string(out[len(head):])
		}
		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
}