Agent-System-Backend/internal/tools/exec_interpreter.go

// run_shell / run_python 解释器执行器（与 register_artifact 同族）。
//
// 语义：模型提交的脚本原样落盘到会话工作区 .agent/ 下（字节精确、执行后保留以便审计），
// 再以该文件为唯一参数启动白名单解释器（bash / python3 的绝对路径）——模型输入不进入
// 进程参数。工作目录锁定会话沙箱；超时取请求 context 与 timeout_sec（默认 60s、上限 300s）
// 的先到者，超时杀掉整个进程组；输出 64KB 头尾截断。非零退出、超时与取消作为文本反馈返回
// （模型据此自纠），仅启动失败算工具错误。
package tools

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"
	"time"

	core "agent-system-core"
)

// RegisterInterpreterTools 注册解释器执行原语（run_shell / run_python）。
func RegisterInterpreterTools(r *Registry) error {
	if err := r.Register(core.ToolDesc{
		Name: "run_shell", Source: core.SourceBuiltin,
		Description: "在会话工作区目录内执行 bash 脚本：检查环境、安装依赖（pip install --user …）、处理文件等。脚本原样落盘执行；返回退出码与合并输出（stdout+stderr，超限头尾截断）。以当前用户身份运行。",
		Parameters: json.RawMessage(`{"type":"object","properties":{` +
			`"script":{"type":"string","description":"bash 脚本（可多行）"},` +
			`"timeout_sec":{"type":"integer","description":"超时秒数（默认 60，上限 300）"}}` +
			`,"required":["script"]}`),
	}, execShell); err != nil {
		return err
	}
	return r.Register(core.ToolDesc{
		Name: "run_python", Source: core.SourceBuiltin,
		Description: "执行 Python 脚本（python3，工作目录为会话工作区）：生成文档/表格/演示稿/数据处理的首选路径。脚本原样落盘执行；stdout/stderr 合并返回，超限头尾截断。python-docx/openpyxl/python-pptx 缺失时可先 pip install --user 安装。",
		Parameters: json.RawMessage(`{"type":"object","properties":{` +
			`"script":{"type":"string","description":"完整 Python 脚本（utf-8）"},` +
			`"timeout_sec":{"type":"integer","description":"超时秒数（默认 60，上限 300）"}}` +
			`,"required":["script"]}`),
	}, execPython)
}

// execEnv 最小环境：PATH/HOME/LANG/TZ/TERM 继承（pip --user、字体、中文 locale），
// 其余变量不带入子进程。
func execEnv() []string {
	var out []string
	for _, kv := range os.Environ() {
		for _, keep := range []string{"PATH=", "HOME=", "LANG=", "LC_", "TZ=", "TERM="} {
			if strings.HasPrefix(kv, keep) {
				out = append(out, kv)
				break
			}
		}
	}
	return append(out, "PYTHONDONTWRITEBYTECODE=1")
}

// truncateOutput 头 2/3 + 尾 1/3 保留，中部标记省略。
func truncateOutput(b []byte) string {
	if len(b) <= execMaxOutput {
		return string(b)
	}
	head := execMaxOutput * 2 / 3
	tail := execMaxOutput / 3
	return string(b[:head]) + fmt.Sprintf("\n…（输出超限，省略 %d 字节）…\n", len(b)-head-tail) + string(b[len(b)-tail:])
}

// writeScript 把模型脚本落盘到工作区 .agent/ 下（字节精确，执行后保留），返回绝对路径。
func writeScript(rc ToolExecCtx, sandbox, prefix, ext, content string) (string, error) {
	dir := filepath.Join(sandbox, execScriptDir)
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return "", err
	}
	p := filepath.Join(dir, fmt.Sprintf("%s-%s-%d%s", prefix, rc.ConversationID(), time.Now().UnixNano(), ext))
	if err := os.WriteFile(p, []byte(content), 0o600); err != nil {
		return "", err
	}
	return p, nil
}

type execContexter interface {
	Context() context.Context
}

func execParentCtx(rc ToolExecCtx) context.Context {
	if c, ok := rc.(execContexter); ok {
		if ctx := c.Context(); ctx != nil {
			return ctx
		}
	}
	return context.Background()
}

// resolveInterpreter 把白名单解释器名钉成绝对路径，避免 PATH 在子进程里被换掉。
func resolveInterpreter(name string) (string, error) {
	switch name {
	case "bash", "python3":
	default:
		return "", fmt.Errorf("interpreter %s is not allowlisted", name)
	}
	p, err := exec.LookPath(name)
	if err != nil {
		return "", err
	}
	if !filepath.IsAbs(p) {
		return "", fmt.Errorf("%s: resolved path is not absolute: %s", name, p)
	}
	return p, nil
}

// runInterpreter 启动白名单解释器执行落盘脚本（脚本路径为唯一参数）。
// 非零退出、超时与请求取消都是正常反馈，以文本返回；仅启动失败算错误。
func runInterpreter(parent context.Context, interpreter, scriptPath, dir string, timeoutS int) (string, error) {
	if parent == nil {
		parent = context.Background()
	}
	if timeoutS <= 0 {
		timeoutS = execDefaultTimeoutS
	}
	if timeoutS > execMaxTimeoutS {
		timeoutS = execMaxTimeoutS
	}
	bin, err := resolveInterpreter(interpreter)
	if err != nil {
		return "", fmt.Errorf("%s: %w", interpreter, err)
	}
	ctx, cancel := context.WithTimeout(parent, time.Duration(timeoutS)*time.Second)
	defer cancel()
	cmd := exec.CommandContext(ctx, bin, scriptPath)
	cmd.Dir = dir
	cmd.Env = execEnv()
	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
	cmd.WaitDelay = 2 * time.Second
	cmd.Cancel = func() error {
		if cmd.Process == nil {
			return os.ErrProcessDone
		}
		err := syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
		if errors.Is(err, syscall.ESRCH) {
			return nil
		}
		return err
	}
	out, err := cmd.CombinedOutput()
	text := truncateOutput(out)
	if errors.Is(ctx.Err(), context.DeadlineExceeded) {
		return fmt.Sprintf("exit=timeout(%ds)\n%s", timeoutS, text), nil
	}
	if errors.Is(parent.Err(), context.Canceled) {
		return fmt.Sprintf("exit=canceled\n%s", text), nil
	}
	if err != nil {
		var ee *exec.ExitError
		if errors.As(err, &ee) {
			return fmt.Sprintf("exit=%d\n%s", ee.ExitCode(), text), nil
		}
		return "", fmt.Errorf("%s: %v：%s", interpreter, err, text)
	}
	return "exit=0\n" + text, nil
}

func execShell(rc ToolExecCtx, args json.RawMessage) (string, error) {
	if err := execGate(rc); err != nil {
		return "", err
	}
	var a struct {
		Script     string `json:"script"`
		Command    string `json:"command"`
		TimeoutS   int    `json:"timeout_sec"`
		TimeoutSec int    `json:"timeoutSec"`
	}
	if len(args) > 0 {
		if err := json.Unmarshal(args, &a); err != nil {
			return "", fmt.Errorf("run_shell: bad args: %w", err)
		}
	}
	script := a.Script
	if strings.TrimSpace(script) == "" {
		script = a.Command
	}
	if strings.TrimSpace(script) == "" {
		return "", fmt.Errorf("run_shell: script required")
	}
	sandbox, err := execSandbox(rc)
	if err != nil {
		return "", err
	}
	p, err := writeScript(rc, sandbox, "sh", ".sh", script)
	if err != nil {
		return "", err
	}
	return runInterpreter(execParentCtx(rc), "bash", p, sandbox, firstPos(a.TimeoutS, a.TimeoutSec))
}

func execPython(rc ToolExecCtx, args json.RawMessage) (string, error) {
	if err := execGate(rc); err != nil {
		return "", err
	}
	var a struct {
		Script     string `json:"script"`
		TimeoutS   int    `json:"timeout_sec"`
		TimeoutSec int    `json:"timeoutSec"`
	}
	if len(args) > 0 {
		if err := json.Unmarshal(args, &a); err != nil {
			return "", fmt.Errorf("run_python: bad args: %w", err)
		}
	}
	if strings.TrimSpace(a.Script) == "" {
		return "", fmt.Errorf("run_python: script required")
	}
	sandbox, err := execSandbox(rc)
	if err != nil {
		return "", err
	}
	p, err := writeScript(rc, sandbox, "py", ".py", a.Script)
	if err != nil {
		return "", err
	}
	return runInterpreter(execParentCtx(rc), "python3", p, sandbox, firstPos(a.TimeoutS, a.TimeoutSec))
}

func firstPos(vals ...int) int {
	for _, v := range vals {
		if v > 0 {
			return v
		}
	}
	return 0
}