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

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

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

	core "agent-system-core"

	agenttools "agent-native/runtime/tools"
)

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

// skillsRoot 由组合根在启动时注入（<data>/skills）。空 = 不注入技能路径。
var skillsRoot string

// SetSkillsRoot 让 run_python / run_shell 能 import housestyle，并看到技能目录。
func SetSkillsRoot(dir string) { skillsRoot = dir }

// execEnv 最小环境：PATH/HOME/LANG/TZ/TERM 继承（pip --user、字体、中文 locale），
// 其余变量不带入子进程。技能根配置后追加 AGENT_SKILLS_DIR 与 PYTHONPATH。
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
			}
		}
	}
	out = append(out, "PYTHONDONTWRITEBYTECODE=1")
	if skillsRoot != "" {
		out = append(out,
			"AGENT_SKILLS_DIR="+skillsRoot,
			"PYTHONPATH="+filepath.Join(skillsRoot, "house-style"),
		)
	}
	return out
}

// 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, sink func(stream, text string)) (string, error) {
	if parent == nil {
		parent = context.Background()
	}
	if timeoutS <= 0 {
		timeoutS = execDefaultTimeoutS
	}
	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
	}
	pr, pw := io.Pipe()
	cmd.Stdout = pw
	cmd.Stderr = pw
	var buf bytes.Buffer
	readDone := make(chan struct{})
	go func() {
		defer close(readDone)
		tmp := make([]byte, 4096)
		for {
			n, rerr := pr.Read(tmp)
			if n > 0 {
				if buf.Len() < 1<<20 {
					buf.Write(tmp[:n])
				}
				if sink != nil {
					sink("stdout", string(tmp[:n]))
				}
			}
			if rerr != nil {
				return
			}
		}
	}()
	err = cmd.Start()
	if err != nil {
		_ = pw.Close()
		<-readDone
		return "", fmt.Errorf("%s: %w", interpreter, err)
	}
	err = cmd.Wait()
	_ = pw.Close()
	<-readDone
	text := truncateOutput(buf.Bytes())
	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), toolOutputSink(rc))
}

// ScriptCommand 准备一条以当前用户身份、工作目录为会话沙箱的 python3 命令。
// 调用方在 Start 之前设置 ExtraFiles、Stdout 和 Stderr。取消 ctx 会杀掉整个进程组。
func ScriptCommand(ctx context.Context, sandbox, path string) (*exec.Cmd, error) {
	if ctx == nil {
		ctx = context.Background()
	}
	bin, err := resolveInterpreter("python3")
	if err != nil {
		return nil, err
	}
	cmd := exec.CommandContext(ctx, bin, path)
	cmd.Dir = sandbox
	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
	}
	return cmd, nil
}

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), toolOutputSink(rc))
}

// toolOutputSink 在执行上下文实现了 EmitToolOutput 时，把 stdout 增量交出去。
func toolOutputSink(rc ToolExecCtx) func(stream, text string) {
	type emitter interface {
		EmitToolOutput(stream, text string)
	}
	e, ok := rc.(emitter)
	if !ok {
		return nil
	}
	return e.EmitToolOutput
}

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

// InterpreterExecConfig 把产品解释器白名单换算为 agent-native 内核的
// ExecConfig（squad 角色经 agent-native 后端执行时复用同一套约束与环境）。
// 两个解释器都缺席时返回 nil（内核侧不注册执行工具，fail-closed）。
func InterpreterExecConfig() *agenttools.ExecConfig {
	cfg := &agenttools.ExecConfig{Env: execEnv()}
	if p, err := exec.LookPath("bash"); err == nil {
		cfg.Bash = p
	}
	if p, err := exec.LookPath("python3"); err == nil {
		cfg.Python = p
	}
	if cfg.Bash == "" && cfg.Python == "" {
		return nil
	}
	return cfg
}