feat(exec): refactor ExecTool to remove bus dependency and improve callback handling
Signed-off-by: Boris Bliznioukov <blib@mail.com>
This commit is contained in:
parent
bf1b07d407
commit
970a56801b
7 changed files with 81 additions and 83 deletions
|
|
@ -168,11 +168,11 @@ func registerSharedTools(
|
||||||
agent.Tools.Register(tools.NewSPITool())
|
agent.Tools.Register(tools.NewSPITool())
|
||||||
}
|
}
|
||||||
|
|
||||||
// Exec tool — created here (not in NewAgentInstance) because it
|
// Exec tool — created here (not in NewAgentInstance) for
|
||||||
// needs bus access for background command result delivery.
|
// consistent tool registration alongside other shared tools.
|
||||||
if cfg.Tools.IsToolEnabled("exec") {
|
if cfg.Tools.IsToolEnabled("exec") {
|
||||||
restrict := cfg.Agents.Defaults.RestrictToWorkspace
|
restrict := cfg.Agents.Defaults.RestrictToWorkspace
|
||||||
execTool, err := tools.NewExecToolWithConfig(agent.Workspace, restrict, cfg, msgBus)
|
execTool, err := tools.NewExecToolWithConfig(agent.Workspace, restrict, cfg)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
logger.ErrorCF("agent", "Failed to create exec tool", map[string]any{"error": err.Error()})
|
logger.ErrorCF("agent", "Failed to create exec tool", map[string]any{"error": err.Error()})
|
||||||
} else {
|
} else {
|
||||||
|
|
|
||||||
|
|
@ -31,7 +31,7 @@ func NewCronTool(
|
||||||
cronService *cron.CronService, executor JobExecutor, msgBus *bus.MessageBus, workspace string, restrict bool,
|
cronService *cron.CronService, executor JobExecutor, msgBus *bus.MessageBus, workspace string, restrict bool,
|
||||||
execTimeout time.Duration, config *config.Config,
|
execTimeout time.Duration, config *config.Config,
|
||||||
) (*CronTool, error) {
|
) (*CronTool, error) {
|
||||||
execTool, err := NewExecToolWithConfig(workspace, restrict, config, nil)
|
execTool, err := NewExecToolWithConfig(workspace, restrict, config)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, fmt.Errorf("unable to configure exec tool: %w", err)
|
return nil, fmt.Errorf("unable to configure exec tool: %w", err)
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -502,6 +502,21 @@ func formatCommand(args []string) string {
|
||||||
return cmd
|
return cmd
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// NormalizeCommandKeys returns a new map with all keys passed through
|
||||||
|
// baseCommand. This ensures user-provided override and modifier keys
|
||||||
|
// match the normalized command names used by ClassifyCommand.
|
||||||
|
// Duplicate keys after normalization are resolved by last-write-wins.
|
||||||
|
func NormalizeCommandKeys[V any](m map[string]V) map[string]V {
|
||||||
|
if len(m) == 0 {
|
||||||
|
return m
|
||||||
|
}
|
||||||
|
normalized := make(map[string]V, len(m))
|
||||||
|
for k, v := range m {
|
||||||
|
normalized[baseCommand(k)] = v
|
||||||
|
}
|
||||||
|
return normalized
|
||||||
|
}
|
||||||
|
|
||||||
// baseCommand extracts the basename from a command path.
|
// baseCommand extracts the basename from a command path.
|
||||||
// On Windows, it additionally lowercases the name and strips known
|
// On Windows, it additionally lowercases the name and strips known
|
||||||
// executable extensions (.exe, .cmd, .bat, .com) so that
|
// executable extensions (.exe, .cmd, .bat, .com) so that
|
||||||
|
|
|
||||||
|
|
@ -162,14 +162,16 @@ func pathAwareExecHandler(env expand.Environ) func(next interp.ExecHandlerFunc)
|
||||||
return next(ctx, args)
|
return next(ctx, args)
|
||||||
}
|
}
|
||||||
|
|
||||||
// Look up command using interpreter's PATH
|
// Look up command using interpreter's PATH.
|
||||||
|
// Hard-fail when the command is not found so that resolution
|
||||||
|
// is always constrained to the sanitized environment. Falling
|
||||||
|
// through to next() would let the default handler resolve via
|
||||||
|
// the process-level PATH, bypassing env restrictions.
|
||||||
path, err := lookPath(env, args[0])
|
path, err := lookPath(env, args[0])
|
||||||
if err != nil {
|
if err != nil {
|
||||||
// Command not found in PATH, try the default handler
|
return fmt.Errorf("%s: %w", args[0], err)
|
||||||
return next(ctx, args)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// Replace command with full path and continue
|
|
||||||
fullArgs := append([]string{path}, args[1:]...)
|
fullArgs := append([]string{path}, args[1:]...)
|
||||||
return next(ctx, fullArgs)
|
return next(ctx, fullArgs)
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -37,7 +37,7 @@ func TestRun_WorkingDir_Windows(t *testing.T) {
|
||||||
Dir: tmpDir,
|
Dir: tmpDir,
|
||||||
Timeout: 5 * time.Second,
|
Timeout: 5 * time.Second,
|
||||||
RiskThreshold: RiskHigh, // cmd.exe is risk=critical
|
RiskThreshold: RiskHigh, // cmd.exe is risk=critical
|
||||||
RiskOverrides: map[string]string{"cmd.exe": "low"},
|
RiskOverrides: map[string]string{"cmd": "low"},
|
||||||
})
|
})
|
||||||
|
|
||||||
if result.IsError {
|
if result.IsError {
|
||||||
|
|
@ -58,7 +58,7 @@ func TestRun_HighThresholdAllowsDel_Windows(t *testing.T) {
|
||||||
Dir: tmpDir,
|
Dir: tmpDir,
|
||||||
Timeout: 5 * time.Second,
|
Timeout: 5 * time.Second,
|
||||||
RiskThreshold: RiskHigh,
|
RiskThreshold: RiskHigh,
|
||||||
RiskOverrides: map[string]string{"cmd.exe": "low"},
|
RiskOverrides: map[string]string{"cmd": "low"},
|
||||||
})
|
})
|
||||||
|
|
||||||
if result.IsError {
|
if result.IsError {
|
||||||
|
|
@ -78,7 +78,7 @@ func TestRun_EnvSanitization_Windows(t *testing.T) {
|
||||||
Dir: t.TempDir(),
|
Dir: t.TempDir(),
|
||||||
Timeout: 5 * time.Second,
|
Timeout: 5 * time.Second,
|
||||||
RiskThreshold: RiskHigh,
|
RiskThreshold: RiskHigh,
|
||||||
RiskOverrides: map[string]string{"cmd.exe": "low"},
|
RiskOverrides: map[string]string{"cmd": "low"},
|
||||||
})
|
})
|
||||||
|
|
||||||
if result.IsError {
|
if result.IsError {
|
||||||
|
|
@ -113,7 +113,7 @@ func TestRun_RiskOverrides_Windows(t *testing.T) {
|
||||||
Dir: t.TempDir(),
|
Dir: t.TempDir(),
|
||||||
Timeout: 5 * time.Second,
|
Timeout: 5 * time.Second,
|
||||||
RiskThreshold: RiskMedium,
|
RiskThreshold: RiskMedium,
|
||||||
RiskOverrides: map[string]string{"cmd.exe": "low"},
|
RiskOverrides: map[string]string{"cmd": "low"},
|
||||||
})
|
})
|
||||||
|
|
||||||
if result.IsError && strings.Contains(result.Output, "blocked") {
|
if result.IsError && strings.Contains(result.Output, "blocked") {
|
||||||
|
|
|
||||||
|
|
@ -6,7 +6,6 @@ import (
|
||||||
"os"
|
"os"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
"github.com/sipeed/picoclaw/pkg/bus"
|
|
||||||
"github.com/sipeed/picoclaw/pkg/config"
|
"github.com/sipeed/picoclaw/pkg/config"
|
||||||
"github.com/sipeed/picoclaw/pkg/tools/shell"
|
"github.com/sipeed/picoclaw/pkg/tools/shell"
|
||||||
"github.com/sipeed/picoclaw/pkg/utils"
|
"github.com/sipeed/picoclaw/pkg/utils"
|
||||||
|
|
@ -22,8 +21,8 @@ var (
|
||||||
// with AST-based risk classification, env sanitization, and file-access sandboxing.
|
// with AST-based risk classification, env sanitization, and file-access sandboxing.
|
||||||
//
|
//
|
||||||
// ExecTool implements AsyncExecutor. When the LLM passes background=true,
|
// ExecTool implements AsyncExecutor. When the LLM passes background=true,
|
||||||
// the command runs in a goroutine and the result is delivered via
|
// the command runs in a goroutine and the result is delivered via the
|
||||||
// bus.PublishInbound (re-entering the agent loop). Without a bus,
|
// AsyncCallback provided by the registry. Without a callback,
|
||||||
// background=true falls back to synchronous execution.
|
// background=true falls back to synchronous execution.
|
||||||
type ExecTool struct {
|
type ExecTool struct {
|
||||||
workingDir string
|
workingDir string
|
||||||
|
|
@ -35,26 +34,22 @@ type ExecTool struct {
|
||||||
argModifiers map[string][]shell.ArgModifier
|
argModifiers map[string][]shell.ArgModifier
|
||||||
envAllowlist []string
|
envAllowlist []string
|
||||||
envSet map[string]string
|
envSet map[string]string
|
||||||
|
|
||||||
bus *bus.MessageBus
|
|
||||||
}
|
}
|
||||||
|
|
||||||
func NewExecTool(workingDir string, restrict bool) (*ExecTool, error) {
|
func NewExecTool(workingDir string, restrict bool) (*ExecTool, error) {
|
||||||
return NewExecToolWithConfig(workingDir, restrict, nil, nil)
|
return NewExecToolWithConfig(workingDir, restrict, nil)
|
||||||
}
|
}
|
||||||
|
|
||||||
func NewExecToolWithConfig(
|
func NewExecToolWithConfig(
|
||||||
workingDir string,
|
workingDir string,
|
||||||
restrict bool,
|
restrict bool,
|
||||||
cfg *config.Config,
|
cfg *config.Config,
|
||||||
msgBus *bus.MessageBus,
|
|
||||||
) (*ExecTool, error) {
|
) (*ExecTool, error) {
|
||||||
t := &ExecTool{
|
t := &ExecTool{
|
||||||
workingDir: workingDir,
|
workingDir: workingDir,
|
||||||
timeout: 60 * time.Second,
|
timeout: 60 * time.Second,
|
||||||
restrictToWorkspace: restrict,
|
restrictToWorkspace: restrict,
|
||||||
riskThreshold: shell.RiskMedium,
|
riskThreshold: shell.RiskMedium,
|
||||||
bus: msgBus,
|
|
||||||
}
|
}
|
||||||
|
|
||||||
if cfg != nil {
|
if cfg != nil {
|
||||||
|
|
@ -70,8 +65,8 @@ func NewExecToolWithConfig(
|
||||||
t.riskThreshold = level
|
t.riskThreshold = level
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
t.riskOverrides = execCfg.RiskOverrides
|
t.riskOverrides = shell.NormalizeCommandKeys(execCfg.RiskOverrides)
|
||||||
t.argModifiers = parseArgModifiers(execCfg.ArgModifiers)
|
t.argModifiers = shell.NormalizeCommandKeys(parseArgModifiers(execCfg.ArgModifiers))
|
||||||
t.envAllowlist = execCfg.EnvAllowlist
|
t.envAllowlist = execCfg.EnvAllowlist
|
||||||
t.envSet = execCfg.EnvSet
|
t.envSet = execCfg.EnvSet
|
||||||
}
|
}
|
||||||
|
|
@ -146,13 +141,12 @@ func (t *ExecTool) Execute(ctx context.Context, args map[string]any) *ToolResult
|
||||||
}
|
}
|
||||||
|
|
||||||
// ExecuteAsync implements AsyncExecutor. The registry calls this for
|
// ExecuteAsync implements AsyncExecutor. The registry calls this for
|
||||||
// parallel tool dispatch. When background=true and a bus is configured,
|
// parallel tool dispatch. When background=true and a callback is provided,
|
||||||
// the command runs in a goroutine and the result is delivered via
|
// the command runs in a goroutine and the result is delivered via cb.
|
||||||
// bus.PublishInbound (re-entering the agent loop). Without a bus,
|
// Without a callback, background=true falls back to synchronous execution.
|
||||||
// background=true falls back to synchronous execution.
|
|
||||||
func (t *ExecTool) ExecuteAsync(ctx context.Context, args map[string]any, cb AsyncCallback) *ToolResult {
|
func (t *ExecTool) ExecuteAsync(ctx context.Context, args map[string]any, cb AsyncCallback) *ToolResult {
|
||||||
background, _ := args["background"].(bool)
|
background, _ := args["background"].(bool)
|
||||||
if !background || t.bus == nil {
|
if !background || cb == nil {
|
||||||
return t.Execute(ctx, args)
|
return t.Execute(ctx, args)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -161,9 +155,6 @@ func (t *ExecTool) ExecuteAsync(ctx context.Context, args map[string]any, cb Asy
|
||||||
return errResult
|
return errResult
|
||||||
}
|
}
|
||||||
|
|
||||||
channel := ToolChannel(ctx)
|
|
||||||
chatID := ToolChatID(ctx)
|
|
||||||
|
|
||||||
go func() {
|
go func() {
|
||||||
result := shell.Run(ctx, cfg)
|
result := shell.Run(ctx, cfg)
|
||||||
|
|
||||||
|
|
@ -174,14 +165,9 @@ func (t *ExecTool) ExecuteAsync(ctx context.Context, args map[string]any, cb Asy
|
||||||
utils.Truncate(cfg.Command, 100), result.Output)
|
utils.Truncate(cfg.Command, 100), result.Output)
|
||||||
}
|
}
|
||||||
|
|
||||||
pubCtx, pubCancel := context.WithTimeout(context.Background(), 5*time.Second)
|
if cb != nil {
|
||||||
defer pubCancel()
|
cb(context.Background(), AsyncResult(content))
|
||||||
_ = t.bus.PublishInbound(pubCtx, bus.InboundMessage{
|
}
|
||||||
Channel: "system",
|
|
||||||
SenderID: fmt.Sprintf("exec:%s", utils.Truncate(cfg.Command, 50)),
|
|
||||||
ChatID: fmt.Sprintf("%s:%s", channel, chatID),
|
|
||||||
Content: content,
|
|
||||||
})
|
|
||||||
}()
|
}()
|
||||||
return AsyncResult(fmt.Sprintf("Running `%s` in background", utils.Truncate(cfg.Command, 100)))
|
return AsyncResult(fmt.Sprintf("Running `%s` in background", utils.Truncate(cfg.Command, 100)))
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -5,9 +5,9 @@ import (
|
||||||
"context"
|
"context"
|
||||||
"os"
|
"os"
|
||||||
"strings"
|
"strings"
|
||||||
|
"sync"
|
||||||
"testing"
|
"testing"
|
||||||
|
|
||||||
"github.com/sipeed/picoclaw/pkg/bus"
|
|
||||||
"github.com/sipeed/picoclaw/pkg/config"
|
"github.com/sipeed/picoclaw/pkg/config"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
@ -29,43 +29,42 @@ func TestExecTool_SyncExecution(t *testing.T) {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestExecTool_BackgroundWithoutBus(t *testing.T) {
|
func TestExecTool_BackgroundWithoutCallback(t *testing.T) {
|
||||||
tool, err := NewExecTool(t.TempDir(), false)
|
tool, err := NewExecTool(t.TempDir(), false)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
|
|
||||||
// background=true but no bus → ExecuteAsync falls back to synchronous
|
// background=true but nil callback → falls back to synchronous
|
||||||
cb := func(_ context.Context, _ *ToolResult) {
|
|
||||||
t.Error("callback should not be invoked for sync fallback")
|
|
||||||
}
|
|
||||||
|
|
||||||
result := tool.ExecuteAsync(context.Background(), map[string]any{
|
result := tool.ExecuteAsync(context.Background(), map[string]any{
|
||||||
"command": "echo fallback",
|
"command": "echo fallback",
|
||||||
"background": true,
|
"background": true,
|
||||||
}, cb)
|
}, nil)
|
||||||
|
|
||||||
if result.Async {
|
if result.Async {
|
||||||
t.Error("should fall back to sync when no bus is configured")
|
t.Error("should fall back to sync when no callback is provided")
|
||||||
}
|
}
|
||||||
if result.IsError {
|
if result.IsError {
|
||||||
t.Fatalf("expected success: %s", result.ForLLM)
|
t.Fatalf("expected success: %s", result.ForLLM)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestExecTool_BackgroundWithBus(t *testing.T) {
|
func TestExecTool_BackgroundWithCallback(t *testing.T) {
|
||||||
msgBus := bus.NewMessageBus()
|
tool, err := NewExecTool(t.TempDir(), false)
|
||||||
defer msgBus.Close()
|
|
||||||
|
|
||||||
tool, err := NewExecToolWithConfig(t.TempDir(), false, nil, msgBus)
|
|
||||||
if err != nil {
|
if err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
|
|
||||||
ctx := WithToolContext(context.Background(), "telegram", "chat-123")
|
var wg sync.WaitGroup
|
||||||
|
wg.Add(1)
|
||||||
|
|
||||||
cb := func(_ context.Context, _ *ToolResult) {}
|
var cbResult *ToolResult
|
||||||
result := tool.ExecuteAsync(ctx, map[string]any{
|
cb := func(_ context.Context, result *ToolResult) {
|
||||||
|
cbResult = result
|
||||||
|
wg.Done()
|
||||||
|
}
|
||||||
|
|
||||||
|
result := tool.ExecuteAsync(context.Background(), map[string]any{
|
||||||
"command": "echo bg_output",
|
"command": "echo bg_output",
|
||||||
"background": true,
|
"background": true,
|
||||||
}, cb)
|
}, cb)
|
||||||
|
|
@ -74,39 +73,35 @@ func TestExecTool_BackgroundWithBus(t *testing.T) {
|
||||||
t.Fatal("expected async result")
|
t.Fatal("expected async result")
|
||||||
}
|
}
|
||||||
|
|
||||||
// Consume the inbound message published by the background goroutine.
|
wg.Wait()
|
||||||
msg, ok := msgBus.ConsumeInbound(ctx)
|
|
||||||
if !ok {
|
|
||||||
t.Fatal("expected inbound message from background exec")
|
|
||||||
}
|
|
||||||
|
|
||||||
if msg.Channel != "system" {
|
if cbResult == nil {
|
||||||
t.Errorf("expected channel=system, got %q", msg.Channel)
|
t.Fatal("callback was not invoked")
|
||||||
}
|
}
|
||||||
if msg.ChatID != "telegram:chat-123" {
|
if !strings.Contains(cbResult.ForLLM, "bg_output") {
|
||||||
t.Errorf("expected chatID=telegram:chat-123, got %q", msg.ChatID)
|
t.Errorf("expected output in callback result: %s", cbResult.ForLLM)
|
||||||
}
|
}
|
||||||
if !strings.Contains(msg.Content, "bg_output") {
|
if !strings.Contains(cbResult.ForLLM, "completed") {
|
||||||
t.Errorf("expected output in message content: %s", msg.Content)
|
t.Errorf("expected 'completed' in callback result: %s", cbResult.ForLLM)
|
||||||
}
|
|
||||||
if !strings.Contains(msg.Content, "completed") {
|
|
||||||
t.Errorf("expected 'completed' in message content: %s", msg.Content)
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestExecTool_BackgroundBlockedCommand(t *testing.T) {
|
func TestExecTool_BackgroundBlockedCommand(t *testing.T) {
|
||||||
msgBus := bus.NewMessageBus()
|
tool, err := NewExecTool(t.TempDir(), false)
|
||||||
defer msgBus.Close()
|
|
||||||
|
|
||||||
tool, err := NewExecToolWithConfig(t.TempDir(), false, nil, msgBus)
|
|
||||||
if err != nil {
|
if err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
|
|
||||||
ctx := WithToolContext(context.Background(), "telegram", "chat-123")
|
var wg sync.WaitGroup
|
||||||
|
wg.Add(1)
|
||||||
|
|
||||||
cb := func(_ context.Context, _ *ToolResult) {}
|
var cbResult *ToolResult
|
||||||
result := tool.ExecuteAsync(ctx, map[string]any{
|
cb := func(_ context.Context, result *ToolResult) {
|
||||||
|
cbResult = result
|
||||||
|
wg.Done()
|
||||||
|
}
|
||||||
|
|
||||||
|
result := tool.ExecuteAsync(context.Background(), map[string]any{
|
||||||
"command": "sudo rm -rf /",
|
"command": "sudo rm -rf /",
|
||||||
"background": true,
|
"background": true,
|
||||||
}, cb)
|
}, cb)
|
||||||
|
|
@ -115,13 +110,13 @@ func TestExecTool_BackgroundBlockedCommand(t *testing.T) {
|
||||||
t.Fatal("expected async result even for blocked commands")
|
t.Fatal("expected async result even for blocked commands")
|
||||||
}
|
}
|
||||||
|
|
||||||
msg, ok := msgBus.ConsumeInbound(ctx)
|
wg.Wait()
|
||||||
if !ok {
|
|
||||||
t.Fatal("expected inbound message from background exec")
|
|
||||||
}
|
|
||||||
|
|
||||||
if !strings.Contains(msg.Content, "failed") {
|
if cbResult == nil {
|
||||||
t.Errorf("expected 'failed' in message content: %s", msg.Content)
|
t.Fatal("callback was not invoked")
|
||||||
|
}
|
||||||
|
if !strings.Contains(cbResult.ForLLM, "failed") {
|
||||||
|
t.Errorf("expected 'failed' in callback result: %s", cbResult.ForLLM)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -245,7 +240,7 @@ func TestNewExecToolWithConfig_EnableDenyPatternsFalseWarning(t *testing.T) {
|
||||||
out := captureStdout(t, func() {
|
out := captureStdout(t, func() {
|
||||||
cfg := &config.Config{}
|
cfg := &config.Config{}
|
||||||
cfg.Tools.Exec.EnableDenyPatterns = boolPtr(false)
|
cfg.Tools.Exec.EnableDenyPatterns = boolPtr(false)
|
||||||
_, err := NewExecToolWithConfig(t.TempDir(), false, cfg, nil)
|
_, err := NewExecToolWithConfig(t.TempDir(), false, cfg)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue