loveuer
117 lines
2.9 KiB
Go
117 lines
2.9 KiB
Go
package tools
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"os"
|
|
"os/exec"
|
|
"strings"
|
|
"sync"
|
|
|
|
"agentu/pkg/llm"
|
|
)
|
|
|
|
type ShellRunTool struct {
|
|
workingDir string
|
|
}
|
|
|
|
func NewShellRunTool(workingDir string) *ShellRunTool {
|
|
return &ShellRunTool{workingDir: workingDir}
|
|
}
|
|
func (*ShellRunTool) Mutates() bool { return true }
|
|
|
|
func (t *ShellRunTool) Definition() llm.Tool {
|
|
return llm.Tool{
|
|
Type: "function",
|
|
Function: llm.ToolFunction{
|
|
Name: "shell_run",
|
|
Description: "Run a non-interactive shell command on the local machine. Use for builds, tests, local status inspection, and local automation. Avoid commands that require an interactive TTY, credentials, or open long-running sessions.",
|
|
Parameters: JSONSchema(`{
|
|
"type": "object",
|
|
"properties": {
|
|
"command": {"type": "string", "description": "Command to execute."},
|
|
"workdir": {"type": "string", "description": "Optional working directory. Relative paths are resolved from the configured agent working directory."}
|
|
},
|
|
"required": ["command"],
|
|
"additionalProperties": false
|
|
}`),
|
|
},
|
|
}
|
|
}
|
|
|
|
type shellRunArgs struct {
|
|
Command string `json:"command"`
|
|
Workdir string `json:"workdir"`
|
|
}
|
|
|
|
func (t *ShellRunTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
|
|
return t.ExecuteStream(ctx, raw, nil)
|
|
}
|
|
|
|
func (t *ShellRunTool) ExecuteStream(ctx context.Context, raw json.RawMessage, emit func(string) error) (string, error) {
|
|
var args shellRunArgs
|
|
if err := json.Unmarshal(raw, &args); err != nil {
|
|
return "", fmt.Errorf("parse arguments: %w", err)
|
|
}
|
|
if strings.TrimSpace(args.Command) == "" {
|
|
return "", errors.New("command is required")
|
|
}
|
|
|
|
workdir, err := resolvePath(t.workingDir, args.Workdir)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
shell := os.Getenv("SHELL")
|
|
if shell == "" {
|
|
shell = "/bin/sh"
|
|
}
|
|
cmd := exec.CommandContext(ctx, shell, "-lc", args.Command)
|
|
cmd.Dir = workdir
|
|
|
|
var mu sync.Mutex
|
|
var combined bytes.Buffer
|
|
writer := &shellStreamWriter{mu: &mu, buf: &combined, emit: emit}
|
|
cmd.Stdout = writer
|
|
cmd.Stderr = writer
|
|
|
|
runErr := cmd.Run()
|
|
text := FormatResult(combined.String())
|
|
if errors.Is(ctx.Err(), context.DeadlineExceeded) {
|
|
return text, fmt.Errorf("command timed out: %w", ctx.Err())
|
|
}
|
|
if runErr != nil {
|
|
if text != "" {
|
|
return text, fmt.Errorf("command failed: %w", runErr)
|
|
}
|
|
return "", fmt.Errorf("command failed: %w", runErr)
|
|
}
|
|
return text, nil
|
|
}
|
|
|
|
// shellStreamWriter captures command output for the final result while
|
|
// forwarding every write to the streaming callback as it happens.
|
|
type shellStreamWriter struct {
|
|
mu *sync.Mutex
|
|
buf *bytes.Buffer
|
|
emit func(string) error
|
|
}
|
|
|
|
func (w *shellStreamWriter) Write(p []byte) (int, error) {
|
|
w.mu.Lock()
|
|
n, writeErr := w.buf.Write(p)
|
|
w.mu.Unlock()
|
|
if writeErr != nil {
|
|
return n, writeErr
|
|
}
|
|
if w.emit != nil {
|
|
if err := w.emit(string(p)); err != nil {
|
|
return 0, err
|
|
}
|
|
}
|
|
return len(p), nil
|
|
}
|