560 lines
18 KiB
Go
560 lines
18 KiB
Go
package agent
|
||
|
||
import (
|
||
"bytes"
|
||
"context"
|
||
"encoding/json"
|
||
"errors"
|
||
"fmt"
|
||
"os"
|
||
"os/exec"
|
||
"path/filepath"
|
||
"sort"
|
||
"strings"
|
||
"time"
|
||
)
|
||
|
||
type Toolset struct {
|
||
Workspace string
|
||
Timeout time.Duration
|
||
MaxOutput int
|
||
shell string
|
||
shellArgs []string
|
||
python string
|
||
dockerSocket func() string
|
||
apiCfg *APIConfigStore
|
||
coop *CoopManager
|
||
// sessionRoot 是各会话独立工作目录的根（data/workspaces），
|
||
// 每个会话在其下拥有 <sessionRoot>/<sessionID> 目录，互不干扰。
|
||
sessionRoot string
|
||
}
|
||
|
||
func NewToolset(cfg Config, dockerCfg *DockerConfigStore, apiCfg *APIConfigStore) *Toolset {
|
||
toolset := &Toolset{
|
||
Workspace: cfg.Workspace,
|
||
Timeout: cfg.ToolTimeout,
|
||
MaxOutput: cfg.MaxToolResultChars,
|
||
apiCfg: apiCfg,
|
||
sessionRoot: filepath.Join(cfg.DataDir, "workspaces"),
|
||
}
|
||
|
||
if dockerCfg != nil {
|
||
toolset.dockerSocket = dockerCfg.Socket
|
||
} else {
|
||
toolset.dockerSocket = defaultDockerSocket
|
||
}
|
||
|
||
if shell, err := exec.LookPath("bash"); err == nil {
|
||
toolset.shell = shell
|
||
toolset.shellArgs = []string{"-lc"}
|
||
} else if shell, err := exec.LookPath("pwsh"); err == nil {
|
||
toolset.shell = shell
|
||
toolset.shellArgs = []string{"-NoProfile", "-NonInteractive", "-Command"}
|
||
} else if shell, err := exec.LookPath("powershell"); err == nil {
|
||
toolset.shell = shell
|
||
toolset.shellArgs = []string{"-NoProfile", "-NonInteractive", "-Command"}
|
||
} else {
|
||
toolset.shell = "cmd.exe"
|
||
toolset.shellArgs = []string{"/C"}
|
||
}
|
||
|
||
if python, err := exec.LookPath("python"); err == nil {
|
||
toolset.python = python
|
||
} else if python, err := exec.LookPath("python3"); err == nil {
|
||
toolset.python = python
|
||
} else {
|
||
toolset.python = "python"
|
||
}
|
||
return toolset
|
||
}
|
||
|
||
func (t *Toolset) Definitions() []ToolDefinition {
|
||
return []ToolDefinition{
|
||
{
|
||
Type: "function",
|
||
Function: ToolDefinitionFunction{
|
||
Name: "run_bash",
|
||
Description: "在用户工作区目录中执行一条 Bash 命令。Windows 上没有 Bash 时自动回退到 PowerShell 或 cmd。适合查看文件、运行构建、安装依赖、搜索代码、启动程序等。",
|
||
Parameters: map[string]any{
|
||
"type": "object",
|
||
"properties": map[string]any{
|
||
"command": map[string]any{
|
||
"type": "string",
|
||
"description": "要执行的完整命令，例如 `ls -la` 或 `go test ./...`。",
|
||
},
|
||
"timeout_seconds": map[string]any{
|
||
"type": "integer",
|
||
"description": "命令超时秒数，默认 120 秒。",
|
||
},
|
||
},
|
||
"required": []string{"command"},
|
||
},
|
||
},
|
||
},
|
||
{
|
||
Type: "function",
|
||
Function: ToolDefinitionFunction{
|
||
Name: "read_file",
|
||
Description: "读取工作区内的文本文件并返回内容。超长文件会被截断。",
|
||
Parameters: map[string]any{
|
||
"type": "object",
|
||
"properties": map[string]any{
|
||
"path": map[string]any{
|
||
"type": "string",
|
||
"description": "文件路径，可以是相对路径或绝对路径。必须位于工作区内。",
|
||
},
|
||
},
|
||
"required": []string{"path"},
|
||
},
|
||
},
|
||
},
|
||
{
|
||
Type: "function",
|
||
Function: ToolDefinitionFunction{
|
||
Name: "write_file",
|
||
Description: "写入或覆盖工作区内的文件。父目录不存在时自动创建。",
|
||
Parameters: map[string]any{
|
||
"type": "object",
|
||
"properties": map[string]any{
|
||
"path": map[string]any{
|
||
"type": "string",
|
||
"description": "目标文件路径，可以是相对路径或绝对路径。必须位于工作区内。",
|
||
},
|
||
"content": map[string]any{
|
||
"type": "string",
|
||
"description": "要写入的完整文件内容。",
|
||
},
|
||
},
|
||
"required": []string{"path", "content"},
|
||
},
|
||
},
|
||
},
|
||
{
|
||
Type: "function",
|
||
Function: ToolDefinitionFunction{
|
||
Name: "run_python",
|
||
Description: "在工作区目录中执行一段 Python 代码。适合数据处理、批量修改、生成脚本和自动化任务。",
|
||
Parameters: map[string]any{
|
||
"type": "object",
|
||
"properties": map[string]any{
|
||
"code": map[string]any{
|
||
"type": "string",
|
||
"description": "要执行的完整 Python 代码。",
|
||
},
|
||
},
|
||
"required": []string{"code"},
|
||
},
|
||
},
|
||
},
|
||
{
|
||
Type: "function",
|
||
Function: ToolDefinitionFunction{
|
||
Name: "list_directory",
|
||
Description: "列出工作区内目录的内容，包括文件大小和修改时间。",
|
||
Parameters: map[string]any{
|
||
"type": "object",
|
||
"properties": map[string]any{
|
||
"path": map[string]any{
|
||
"type": "string",
|
||
"description": "目录路径，默认为工作区根目录。",
|
||
},
|
||
},
|
||
},
|
||
},
|
||
},
|
||
{
|
||
Type: "function",
|
||
Function: ToolDefinitionFunction{
|
||
Name: "docker_ps",
|
||
Description: "列出 Docker 容器。默认连接本地 Docker，可通过 Web 设置中的 Docker Socket 切换到远程 Docker 服务器。",
|
||
Parameters: map[string]any{
|
||
"type": "object",
|
||
"properties": map[string]any{
|
||
"all": map[string]any{
|
||
"type": "boolean",
|
||
"description": "是否列出所有容器（包括已停止的），默认 false 只显示运行中。",
|
||
},
|
||
},
|
||
},
|
||
},
|
||
},
|
||
{
|
||
Type: "function",
|
||
Function: ToolDefinitionFunction{
|
||
Name: "docker_images",
|
||
Description: "列出 Docker 镜像列表。",
|
||
Parameters: map[string]any{
|
||
"type": "object",
|
||
"properties": map[string]any{},
|
||
},
|
||
},
|
||
},
|
||
{
|
||
Type: "function",
|
||
Function: ToolDefinitionFunction{
|
||
Name: "docker_logs",
|
||
Description: "查看 Docker 容器的最近日志。",
|
||
Parameters: map[string]any{
|
||
"type": "object",
|
||
"properties": map[string]any{
|
||
"container": map[string]any{
|
||
"type": "string",
|
||
"description": "容器名称或 ID。",
|
||
},
|
||
"tail": map[string]any{
|
||
"type": "integer",
|
||
"description": "返回最近的日志行数，默认 100。",
|
||
},
|
||
},
|
||
"required": []string{"container"},
|
||
},
|
||
},
|
||
},
|
||
{
|
||
Type: "function",
|
||
Function: ToolDefinitionFunction{
|
||
Name: "docker_inspect",
|
||
Description: "查看 Docker 容器的详细配置和状态信息。",
|
||
Parameters: map[string]any{
|
||
"type": "object",
|
||
"properties": map[string]any{
|
||
"container": map[string]any{
|
||
"type": "string",
|
||
"description": "容器名称或 ID。",
|
||
},
|
||
},
|
||
"required": []string{"container"},
|
||
},
|
||
},
|
||
},
|
||
{
|
||
Type: "function",
|
||
Function: ToolDefinitionFunction{
|
||
Name: "docker_exec",
|
||
Description: "在 Docker 容器内执行一条命令（使用 sh -c），返回执行输出。",
|
||
Parameters: map[string]any{
|
||
"type": "object",
|
||
"properties": map[string]any{
|
||
"container": map[string]any{
|
||
"type": "string",
|
||
"description": "容器名称或 ID。",
|
||
},
|
||
"command": map[string]any{
|
||
"type": "string",
|
||
"description": "要在容器内执行的命令，例如 `ls -la` 或 `cat /etc/os-release`。",
|
||
},
|
||
},
|
||
"required": []string{"container", "command"},
|
||
},
|
||
},
|
||
},
|
||
{
|
||
Type: "function",
|
||
Function: ToolDefinitionFunction{
|
||
Name: "docker_start",
|
||
Description: "启动一个 Docker 容器。",
|
||
Parameters: map[string]any{
|
||
"type": "object",
|
||
"properties": map[string]any{
|
||
"container": map[string]any{
|
||
"type": "string",
|
||
"description": "容器名称或 ID。",
|
||
},
|
||
},
|
||
"required": []string{"container"},
|
||
},
|
||
},
|
||
},
|
||
{
|
||
Type: "function",
|
||
Function: ToolDefinitionFunction{
|
||
Name: "docker_stop",
|
||
Description: "停止一个 Docker 容器。",
|
||
Parameters: map[string]any{
|
||
"type": "object",
|
||
"properties": map[string]any{
|
||
"container": map[string]any{
|
||
"type": "string",
|
||
"description": "容器名称或 ID。",
|
||
},
|
||
"timeout": map[string]any{
|
||
"type": "integer",
|
||
"description": "等待优雅停止的秒数，默认由 Docker 决定。",
|
||
},
|
||
},
|
||
"required": []string{"container"},
|
||
},
|
||
},
|
||
},
|
||
{
|
||
Type: "function",
|
||
Function: ToolDefinitionFunction{
|
||
Name: "run_coop",
|
||
Description: "调用协作容器异步完成一个任务（容器内单个 agent 独立执行，自动注入当前 Agent 已配置的模型、Base URL、API Key 与接口协议，无需重复填写）。支持三种引擎：engine=pi（默认，pi-coding-agent，支持 openai+anthropic 协议）、engine=pigo（pigo，支持 openai+anthropic 协议）或 engine=claude（Claude Code，仅支持 anthropic 协议）。需先构建对应镜像（pi-coop / pigo-coop / claude-coop）。该工具是异步的：调用后立即返回任务 ID，容器在后台运行；任务完成（成功 / 失败 / 超时）后系统会自动注入一条【协作任务完成通知】消息（含结构化结果摘要），由你确认结果、必要时提交 flag 并向用户汇报，无需在此等待。",
|
||
Parameters: map[string]any{
|
||
"type": "object",
|
||
"properties": map[string]any{
|
||
"task": map[string]any{
|
||
"type": "string",
|
||
"description": "要交给协作容器完成的任务描述，例如“请为 XX 编写设计文档并实现原型”。解题类任务需写明目标地址、平台提交规则（BENCHMARK_TOKEN / unique_code / submit API）。",
|
||
},
|
||
"engine": map[string]any{
|
||
"type": "string",
|
||
"description": "协作引擎：pi（默认，pi-coding-agent）、pigo（pigo）或 claude（Claude Code，仅支持 anthropic 协议）。pi 与 pigo 支持 openai+anthropic 协议；claude 仅支持 anthropic 协议。",
|
||
"enum": []string{"pi", "pigo", "claude"},
|
||
},
|
||
"round_max": map[string]any{
|
||
"type": "integer",
|
||
"description": "最大轮次，默认 1（单 agent 一次运行完成，未完成则失败并重新下发）。",
|
||
},
|
||
"timeout": map[string]any{
|
||
"type": "integer",
|
||
"description": "单轮超时秒数，默认 1800（30 分钟）。单 agent 需在此时限内完成全部工作并创建 DONE。",
|
||
},
|
||
"blackboard": map[string]any{
|
||
"type": "string",
|
||
"description": "可选：主机上的目录路径，挂载到容器 /blackboard 保留黑板产物与会话；Windows 路径会自动转换为 WSL 挂载路径（/mnt/盘符/...）。留空则自动使用本会话工作区下的独立子目录 blackboard/<任务ID>（同一会话多次协作互不污染）。",
|
||
},
|
||
},
|
||
"required": []string{"task"},
|
||
},
|
||
},
|
||
},
|
||
}
|
||
}
|
||
|
||
func (t *Toolset) Execute(ctx context.Context, sessionID, name, arguments string) ToolResult {
|
||
var params map[string]any
|
||
if err := json.Unmarshal([]byte(arguments), ¶ms); err != nil {
|
||
return ToolResult{Success: false, Output: "无法解析工具参数: " + err.Error()}
|
||
}
|
||
|
||
switch name {
|
||
case "run_bash":
|
||
command, _ := params["command"].(string)
|
||
if strings.TrimSpace(command) == "" {
|
||
return ToolResult{Success: false, Output: "command 不能为空"}
|
||
}
|
||
return t.runShell(ctx, sessionID, command, intParam(params, "timeout_seconds"))
|
||
case "read_file":
|
||
path, _ := params["path"].(string)
|
||
return t.readFile(sessionID, path)
|
||
case "write_file":
|
||
path, _ := params["path"].(string)
|
||
content, _ := params["content"].(string)
|
||
return t.writeFile(sessionID, path, content)
|
||
case "run_python":
|
||
code, _ := params["code"].(string)
|
||
if strings.TrimSpace(code) == "" {
|
||
return ToolResult{Success: false, Output: "code 不能为空"}
|
||
}
|
||
return t.runPython(ctx, sessionID, code)
|
||
case "list_directory":
|
||
path, _ := params["path"].(string)
|
||
return t.listDirectory(sessionID, path)
|
||
case "docker_ps":
|
||
return t.dockerPS(ctx, params)
|
||
case "docker_images":
|
||
return t.dockerImages(ctx, params)
|
||
case "docker_logs":
|
||
return t.dockerLogs(ctx, params)
|
||
case "docker_inspect":
|
||
return t.dockerInspect(ctx, params)
|
||
case "docker_exec":
|
||
return t.dockerExec(ctx, params)
|
||
case "docker_start":
|
||
return t.dockerStart(ctx, params)
|
||
case "docker_stop":
|
||
return t.dockerStop(ctx, params)
|
||
case "run_coop":
|
||
return t.runCoop(ctx, sessionID, params)
|
||
default:
|
||
return ToolResult{Success: false, Output: "未知工具: " + name}
|
||
}
|
||
}
|
||
|
||
func (t *Toolset) runShell(ctx context.Context, sessionID, command string, timeoutSeconds int) ToolResult {
|
||
// 不信任 LLM 传入的超长 timeout：超过全局工具超时（默认 120s）一律截断为全局值。
|
||
// 否则 sleep / 轮询类命令会长时间阻塞主 Agent（持有会话锁），期间前端
|
||
// 长时间收不到事件会被判"任务中断"，且 coop 完成通知也拿不到锁无法及时处理。
|
||
if timeoutSeconds <= 0 || timeoutSeconds > int(t.Timeout.Seconds()) {
|
||
timeoutSeconds = int(t.Timeout.Seconds())
|
||
}
|
||
ctx, cancel := context.WithTimeout(ctx, time.Duration(timeoutSeconds)*time.Second)
|
||
defer cancel()
|
||
|
||
var stdout, stderr bytes.Buffer
|
||
cmd := exec.CommandContext(ctx, t.shell, t.shellArgs...)
|
||
cmd.Args = append(cmd.Args, command)
|
||
cmd.Dir = t.workspaceFor(sessionID)
|
||
cmd.Env = append(os.Environ(), "AGENT_WORKSPACE="+t.workspaceFor(sessionID))
|
||
cmd.Stdout = &stdout
|
||
cmd.Stderr = &stderr
|
||
|
||
err := cmd.Run()
|
||
var out strings.Builder
|
||
out.WriteString(strings.TrimSpace(stdout.String()))
|
||
if stderrText := strings.TrimSpace(stderr.String()); stderrText != "" {
|
||
if out.Len() > 0 {
|
||
out.WriteString("\n")
|
||
}
|
||
out.WriteString("[stderr]\n")
|
||
out.WriteString(stderrText)
|
||
}
|
||
if err != nil {
|
||
if out.Len() > 0 {
|
||
out.WriteString("\n")
|
||
}
|
||
out.WriteString("[命令失败] " + err.Error())
|
||
}
|
||
return ToolResult{Success: err == nil, Output: t.truncate(out.String())}
|
||
}
|
||
|
||
func (t *Toolset) runPython(ctx context.Context, sessionID, code string) ToolResult {
|
||
ctx, cancel := context.WithTimeout(ctx, t.Timeout)
|
||
defer cancel()
|
||
|
||
var stdout, stderr bytes.Buffer
|
||
cmd := exec.CommandContext(ctx, t.python, "-c", code)
|
||
cmd.Dir = t.workspaceFor(sessionID)
|
||
cmd.Env = append(os.Environ(), "AGENT_WORKSPACE="+t.workspaceFor(sessionID))
|
||
cmd.Stdout = &stdout
|
||
cmd.Stderr = &stderr
|
||
|
||
err := cmd.Run()
|
||
var out strings.Builder
|
||
out.WriteString(strings.TrimSpace(stdout.String()))
|
||
if stderrText := strings.TrimSpace(stderr.String()); stderrText != "" {
|
||
if out.Len() > 0 {
|
||
out.WriteString("\n")
|
||
}
|
||
out.WriteString("[stderr]\n")
|
||
out.WriteString(stderrText)
|
||
}
|
||
if err != nil {
|
||
if out.Len() > 0 {
|
||
out.WriteString("\n")
|
||
}
|
||
out.WriteString("[运行失败] " + err.Error())
|
||
}
|
||
return ToolResult{Success: err == nil, Output: t.truncate(out.String())}
|
||
}
|
||
|
||
func (t *Toolset) readFile(sessionID, path string) ToolResult {
|
||
absPath, err := t.resolvePath(sessionID, path)
|
||
if err != nil {
|
||
return ToolResult{Success: false, Output: err.Error()}
|
||
}
|
||
data, err := os.ReadFile(absPath)
|
||
if err != nil {
|
||
return ToolResult{Success: false, Output: err.Error()}
|
||
}
|
||
return ToolResult{Success: true, Output: t.truncate(string(data))}
|
||
}
|
||
|
||
func (t *Toolset) writeFile(sessionID, path, content string) ToolResult {
|
||
absPath, err := t.resolvePath(sessionID, path)
|
||
if err != nil {
|
||
return ToolResult{Success: false, Output: err.Error()}
|
||
}
|
||
if err := os.MkdirAll(filepath.Dir(absPath), 0o755); err != nil {
|
||
return ToolResult{Success: false, Output: err.Error()}
|
||
}
|
||
if err := os.WriteFile(absPath, []byte(content), 0o644); err != nil {
|
||
return ToolResult{Success: false, Output: err.Error()}
|
||
}
|
||
return ToolResult{Success: true, Output: fmt.Sprintf("已写入 %d 字节到 %s", len(content), absPath)}
|
||
}
|
||
|
||
func (t *Toolset) listDirectory(sessionID, path string) ToolResult {
|
||
absPath, err := t.resolvePath(sessionID, path)
|
||
if err != nil {
|
||
return ToolResult{Success: false, Output: err.Error()}
|
||
}
|
||
entries, err := os.ReadDir(absPath)
|
||
if err != nil {
|
||
return ToolResult{Success: false, Output: err.Error()}
|
||
}
|
||
sort.Slice(entries, func(i, j int) bool {
|
||
return entries[i].Name() < entries[j].Name()
|
||
})
|
||
|
||
var out strings.Builder
|
||
for _, entry := range entries {
|
||
info, infoErr := entry.Info()
|
||
if infoErr != nil {
|
||
continue
|
||
}
|
||
name := entry.Name()
|
||
if entry.IsDir() {
|
||
name += "/"
|
||
}
|
||
fmt.Fprintf(&out, "%-40s %10d %s\n", name, info.Size(), info.ModTime().Format("2006-01-02 15:04"))
|
||
}
|
||
if out.Len() == 0 {
|
||
return ToolResult{Success: true, Output: "目录为空"}
|
||
}
|
||
return ToolResult{Success: true, Output: t.truncate(out.String())}
|
||
}
|
||
|
||
// resolvePath 把工具传入的路径解析为绝对路径：
|
||
// - 相对路径基于该会话的独立工作目录（data/workspaces/<sessionID>）；
|
||
// - 绝对路径允许在项目根目录内（协作黑板产物、agent 自身文件等共享内容）；
|
||
// - 项目根目录之外一律拒绝。
|
||
func (t *Toolset) resolvePath(sessionID, raw string) (string, error) {
|
||
if strings.TrimSpace(raw) == "" {
|
||
raw = "."
|
||
}
|
||
path := filepath.FromSlash(raw)
|
||
if !filepath.IsAbs(path) {
|
||
path = filepath.Join(t.workspaceFor(sessionID), path)
|
||
}
|
||
absPath, err := filepath.Abs(path)
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
if !pathWithin(absPath, t.Workspace) {
|
||
return "", errors.New("路径超出项目根目录范围，已拒绝: " + absPath)
|
||
}
|
||
return absPath, nil
|
||
}
|
||
|
||
// workspaceFor 返回某会话的独立工作目录并确保其存在。
|
||
func (t *Toolset) workspaceFor(sessionID string) string {
|
||
root := t.sessionRoot
|
||
if root == "" {
|
||
root = t.Workspace
|
||
}
|
||
dir := filepath.Join(root, sessionID)
|
||
_ = os.MkdirAll(dir, 0o755)
|
||
return dir
|
||
}
|
||
|
||
// pathWithin 判断 child 是否位于 parent 目录内（含自身）。
|
||
func pathWithin(child, parent string) bool {
|
||
rel, err := filepath.Rel(parent, child)
|
||
if err != nil {
|
||
return false
|
||
}
|
||
return rel == "." || (rel != ".." && !strings.HasPrefix(rel, ".."+string(filepath.Separator)))
|
||
}
|
||
|
||
func (t *Toolset) truncate(value string) string {
|
||
runes := []rune(value)
|
||
if len(runes) <= t.MaxOutput {
|
||
return value
|
||
}
|
||
return string(runes[:t.MaxOutput]) + "\n...[输出过长，已截断]"
|
||
}
|
||
|
||
func intParam(params map[string]any, key string) int {
|
||
switch value := params[key].(type) {
|
||
case float64:
|
||
return int(value)
|
||
case int:
|
||
return value
|
||
default:
|
||
return 0
|
||
}
|
||
}
|