diff --git a/go.mod b/go.mod index 7413646a4..055bc6948 100644 --- a/go.mod +++ b/go.mod @@ -27,6 +27,7 @@ require ( github.com/spf13/cobra v1.10.2 github.com/stretchr/testify v1.11.1 github.com/tencent-connect/botgo v0.2.1 + github.com/traefik/yaegi v0.16.1 go.mau.fi/whatsmeow v0.0.0-20260219150138-7ae702b1eed4 golang.org/x/oauth2 v0.35.0 golang.org/x/time v0.14.0 diff --git a/pkg/tools/go_eval.go b/pkg/tools/go_eval.go index 5a8547d02..eaf730427 100644 --- a/pkg/tools/go_eval.go +++ b/pkg/tools/go_eval.go @@ -1,15 +1,39 @@ package tools import ( - "bytes" "context" "fmt" - "os" - "os/exec" - "path/filepath" + "sync" "time" + + "github.com/traefik/yaegi/interp" + "github.com/traefik/yaegi/stdlib" ) +type threadSafeBuffer struct { + b []byte + mu sync.Mutex +} + +func (b *threadSafeBuffer) Write(p []byte) (n int, err error) { + b.mu.Lock() + defer b.mu.Unlock() + b.b = append(b.b, p...) + return len(p), nil +} + +func (b *threadSafeBuffer) String() string { + b.mu.Lock() + defer b.mu.Unlock() + return string(b.b) +} + +func (b *threadSafeBuffer) Len() int { + b.mu.Lock() + defer b.mu.Unlock() + return len(b.b) +} + type GoEvalTool struct { workspace string timeout time.Duration @@ -27,7 +51,7 @@ func (t *GoEvalTool) Name() string { } func (t *GoEvalTool) Description() string { - return "Executes Go code dynamically. Provide valid Go source code containing a 'main' function. The code will be saved to a temporary file, compiled, and executed. Useful for complex logic or tasks that require writing a Go script." + return "Executes Go code dynamically using Yaegi interpreter. Provide valid Go source code. The code will be interpreted and executed safely without requiring the Go toolchain. Useful for complex logic or tasks that require writing a Go script." } func (t *GoEvalTool) Parameters() map[string]any { @@ -36,7 +60,7 @@ func (t *GoEvalTool) Parameters() map[string]any { "properties": map[string]any{ "code": map[string]any{ "type": "string", - "description": "Valid Go source code to execute (must include 'package main' and 'func main()').", + "description": "Valid Go source code to execute. It does not need to be a complete package with main func, scripts can just be valid Go statements or functions.", }, }, "required": []string{"code"}, @@ -49,46 +73,38 @@ func (t *GoEvalTool) Execute(ctx context.Context, args map[string]any) *ToolResu return ErrorResult("code is required") } - // Create a temporary directory for the Go code - tmpDir, err := os.MkdirTemp(t.workspace, "go_eval_*") - if err != nil { - return ErrorResult(fmt.Sprintf("failed to create temp dir: %v", err)) - } - defer os.RemoveAll(tmpDir) - - // Write code to main.go - mainFile := filepath.Join(tmpDir, "main.go") - if err := os.WriteFile(mainFile, []byte(code), 0600); err != nil { - return ErrorResult(fmt.Sprintf("failed to write code to file: %v", err)) - } - - // Initialize a go module to allow imports from standard library - cmdMod := exec.Command("go", "mod", "init", "goeval") - cmdMod.Dir = tmpDir - if out, err := cmdMod.CombinedOutput(); err != nil { - return ErrorResult(fmt.Sprintf("failed to init go mod: %v\nOutput: %s", err, string(out))) - } - - // Tidy the module to fetch any dependencies - cmdTidy := exec.Command("go", "mod", "tidy") - cmdTidy.Dir = tmpDir - if out, err := cmdTidy.CombinedOutput(); err != nil { - // Log the error but don't fail immediately, it might just be standard library - _ = out - } - - // Run the code cmdCtx, cancel := context.WithTimeout(ctx, t.timeout) defer cancel() - cmdRun := exec.CommandContext(cmdCtx, "go", "run", "main.go") - cmdRun.Dir = tmpDir + var stdout, stderr threadSafeBuffer - var stdout, stderr bytes.Buffer - cmdRun.Stdout = &stdout - cmdRun.Stderr = &stderr + i := interp.New(interp.Options{ + Stdout: &stdout, + Stderr: &stderr, + Args: []string{}, + }) - err = cmdRun.Run() + if err := i.Use(stdlib.Symbols); err != nil { + return &ToolResult{ + ForLLM: fmt.Sprintf("Failed to initialize standard library symbols: %v", err), + ForUser: "Execution environment setup failed.", + IsError: true, + } + } + + // Channel to capture execution result + done := make(chan error, 1) + go func() { + _, err := i.EvalWithContext(cmdCtx, code) + done <- err + }() + + var err error + select { + case <-cmdCtx.Done(): + err = cmdCtx.Err() + case err = <-done: + } output := stdout.String() if stderr.Len() > 0 { @@ -99,7 +115,7 @@ func (t *GoEvalTool) Execute(ctx context.Context, args map[string]any) *ToolResu } if err != nil { - if cmdCtx.Err() == context.DeadlineExceeded { + if err == context.DeadlineExceeded { return &ToolResult{ ForLLM: fmt.Sprintf("Execution timed out after %v.\nOutput so far:\n%s", t.timeout, output), ForUser: "Execution timed out.",