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())
|
||||
}
|
||||
|
||||
// Exec tool — created here (not in NewAgentInstance) because it
|
||||
// needs bus access for background command result delivery.
|
||||
// Exec tool — created here (not in NewAgentInstance) for
|
||||
// consistent tool registration alongside other shared tools.
|
||||
if cfg.Tools.IsToolEnabled("exec") {
|
||||
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 {
|
||||
logger.ErrorCF("agent", "Failed to create exec tool", map[string]any{"error": err.Error()})
|
||||
} else {
|
||||
|
|
|
|||
|
|
@ -31,7 +31,7 @@ func NewCronTool(
|
|||
cronService *cron.CronService, executor JobExecutor, msgBus *bus.MessageBus, workspace string, restrict bool,
|
||||
execTimeout time.Duration, config *config.Config,
|
||||
) (*CronTool, error) {
|
||||
execTool, err := NewExecToolWithConfig(workspace, restrict, config, nil)
|
||||
execTool, err := NewExecToolWithConfig(workspace, restrict, config)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to configure exec tool: %w", err)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -502,6 +502,21 @@ func formatCommand(args []string) string {
|
|||
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.
|
||||
// On Windows, it additionally lowercases the name and strips known
|
||||
// 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)
|
||||
}
|
||||
|
||||
// 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])
|
||||
if err != nil {
|
||||
// Command not found in PATH, try the default handler
|
||||
return next(ctx, args)
|
||||
return fmt.Errorf("%s: %w", args[0], err)
|
||||
}
|
||||
|
||||
// Replace command with full path and continue
|
||||
fullArgs := append([]string{path}, args[1:]...)
|
||||
return next(ctx, fullArgs)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -37,7 +37,7 @@ func TestRun_WorkingDir_Windows(t *testing.T) {
|
|||
Dir: tmpDir,
|
||||
Timeout: 5 * time.Second,
|
||||
RiskThreshold: RiskHigh, // cmd.exe is risk=critical
|
||||
RiskOverrides: map[string]string{"cmd.exe": "low"},
|
||||
RiskOverrides: map[string]string{"cmd": "low"},
|
||||
})
|
||||
|
||||
if result.IsError {
|
||||
|
|
@ -58,7 +58,7 @@ func TestRun_HighThresholdAllowsDel_Windows(t *testing.T) {
|
|||
Dir: tmpDir,
|
||||
Timeout: 5 * time.Second,
|
||||
RiskThreshold: RiskHigh,
|
||||
RiskOverrides: map[string]string{"cmd.exe": "low"},
|
||||
RiskOverrides: map[string]string{"cmd": "low"},
|
||||
})
|
||||
|
||||
if result.IsError {
|
||||
|
|
@ -78,7 +78,7 @@ func TestRun_EnvSanitization_Windows(t *testing.T) {
|
|||
Dir: t.TempDir(),
|
||||
Timeout: 5 * time.Second,
|
||||
RiskThreshold: RiskHigh,
|
||||
RiskOverrides: map[string]string{"cmd.exe": "low"},
|
||||
RiskOverrides: map[string]string{"cmd": "low"},
|
||||
})
|
||||
|
||||
if result.IsError {
|
||||
|
|
@ -113,7 +113,7 @@ func TestRun_RiskOverrides_Windows(t *testing.T) {
|
|||
Dir: t.TempDir(),
|
||||
Timeout: 5 * time.Second,
|
||||
RiskThreshold: RiskMedium,
|
||||
RiskOverrides: map[string]string{"cmd.exe": "low"},
|
||||
RiskOverrides: map[string]string{"cmd": "low"},
|
||||
})
|
||||
|
||||
if result.IsError && strings.Contains(result.Output, "blocked") {
|
||||
|
|
|
|||
|
|
@ -6,7 +6,6 @@ import (
|
|||
"os"
|
||||
"time"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/bus"
|
||||
"github.com/sipeed/picoclaw/pkg/config"
|
||||
"github.com/sipeed/picoclaw/pkg/tools/shell"
|
||||
"github.com/sipeed/picoclaw/pkg/utils"
|
||||
|
|
@ -22,8 +21,8 @@ var (
|
|||
// with AST-based risk classification, env sanitization, and file-access sandboxing.
|
||||
//
|
||||
// ExecTool implements AsyncExecutor. When the LLM passes background=true,
|
||||
// the command runs in a goroutine and the result is delivered via
|
||||
// bus.PublishInbound (re-entering the agent loop). Without a bus,
|
||||
// the command runs in a goroutine and the result is delivered via the
|
||||
// AsyncCallback provided by the registry. Without a callback,
|
||||
// background=true falls back to synchronous execution.
|
||||
type ExecTool struct {
|
||||
workingDir string
|
||||
|
|
@ -35,26 +34,22 @@ type ExecTool struct {
|
|||
argModifiers map[string][]shell.ArgModifier
|
||||
envAllowlist []string
|
||||
envSet map[string]string
|
||||
|
||||
bus *bus.MessageBus
|
||||
}
|
||||
|
||||
func NewExecTool(workingDir string, restrict bool) (*ExecTool, error) {
|
||||
return NewExecToolWithConfig(workingDir, restrict, nil, nil)
|
||||
return NewExecToolWithConfig(workingDir, restrict, nil)
|
||||
}
|
||||
|
||||
func NewExecToolWithConfig(
|
||||
workingDir string,
|
||||
restrict bool,
|
||||
cfg *config.Config,
|
||||
msgBus *bus.MessageBus,
|
||||
) (*ExecTool, error) {
|
||||
t := &ExecTool{
|
||||
workingDir: workingDir,
|
||||
timeout: 60 * time.Second,
|
||||
restrictToWorkspace: restrict,
|
||||
riskThreshold: shell.RiskMedium,
|
||||
bus: msgBus,
|
||||
}
|
||||
|
||||
if cfg != nil {
|
||||
|
|
@ -70,8 +65,8 @@ func NewExecToolWithConfig(
|
|||
t.riskThreshold = level
|
||||
}
|
||||
}
|
||||
t.riskOverrides = execCfg.RiskOverrides
|
||||
t.argModifiers = parseArgModifiers(execCfg.ArgModifiers)
|
||||
t.riskOverrides = shell.NormalizeCommandKeys(execCfg.RiskOverrides)
|
||||
t.argModifiers = shell.NormalizeCommandKeys(parseArgModifiers(execCfg.ArgModifiers))
|
||||
t.envAllowlist = execCfg.EnvAllowlist
|
||||
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
|
||||
// parallel tool dispatch. When background=true and a bus is configured,
|
||||
// the command runs in a goroutine and the result is delivered via
|
||||
// bus.PublishInbound (re-entering the agent loop). Without a bus,
|
||||
// background=true falls back to synchronous execution.
|
||||
// parallel tool dispatch. When background=true and a callback is provided,
|
||||
// the command runs in a goroutine and the result is delivered via cb.
|
||||
// Without a callback, background=true falls back to synchronous execution.
|
||||
func (t *ExecTool) ExecuteAsync(ctx context.Context, args map[string]any, cb AsyncCallback) *ToolResult {
|
||||
background, _ := args["background"].(bool)
|
||||
if !background || t.bus == nil {
|
||||
if !background || cb == nil {
|
||||
return t.Execute(ctx, args)
|
||||
}
|
||||
|
||||
|
|
@ -161,9 +155,6 @@ func (t *ExecTool) ExecuteAsync(ctx context.Context, args map[string]any, cb Asy
|
|||
return errResult
|
||||
}
|
||||
|
||||
channel := ToolChannel(ctx)
|
||||
chatID := ToolChatID(ctx)
|
||||
|
||||
go func() {
|
||||
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)
|
||||
}
|
||||
|
||||
pubCtx, pubCancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer pubCancel()
|
||||
_ = 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,
|
||||
})
|
||||
if cb != nil {
|
||||
cb(context.Background(), AsyncResult(content))
|
||||
}
|
||||
}()
|
||||
return AsyncResult(fmt.Sprintf("Running `%s` in background", utils.Truncate(cfg.Command, 100)))
|
||||
}
|
||||
|
|
|
|||
|
|
@ -5,9 +5,9 @@ import (
|
|||
"context"
|
||||
"os"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/bus"
|
||||
"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)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// background=true but no bus → ExecuteAsync falls back to synchronous
|
||||
cb := func(_ context.Context, _ *ToolResult) {
|
||||
t.Error("callback should not be invoked for sync fallback")
|
||||
}
|
||||
|
||||
// background=true but nil callback → falls back to synchronous
|
||||
result := tool.ExecuteAsync(context.Background(), map[string]any{
|
||||
"command": "echo fallback",
|
||||
"background": true,
|
||||
}, cb)
|
||||
}, nil)
|
||||
|
||||
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 {
|
||||
t.Fatalf("expected success: %s", result.ForLLM)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecTool_BackgroundWithBus(t *testing.T) {
|
||||
msgBus := bus.NewMessageBus()
|
||||
defer msgBus.Close()
|
||||
|
||||
tool, err := NewExecToolWithConfig(t.TempDir(), false, nil, msgBus)
|
||||
func TestExecTool_BackgroundWithCallback(t *testing.T) {
|
||||
tool, err := NewExecTool(t.TempDir(), false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
ctx := WithToolContext(context.Background(), "telegram", "chat-123")
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(1)
|
||||
|
||||
cb := func(_ context.Context, _ *ToolResult) {}
|
||||
result := tool.ExecuteAsync(ctx, map[string]any{
|
||||
var cbResult *ToolResult
|
||||
cb := func(_ context.Context, result *ToolResult) {
|
||||
cbResult = result
|
||||
wg.Done()
|
||||
}
|
||||
|
||||
result := tool.ExecuteAsync(context.Background(), map[string]any{
|
||||
"command": "echo bg_output",
|
||||
"background": true,
|
||||
}, cb)
|
||||
|
|
@ -74,39 +73,35 @@ func TestExecTool_BackgroundWithBus(t *testing.T) {
|
|||
t.Fatal("expected async result")
|
||||
}
|
||||
|
||||
// Consume the inbound message published by the background goroutine.
|
||||
msg, ok := msgBus.ConsumeInbound(ctx)
|
||||
if !ok {
|
||||
t.Fatal("expected inbound message from background exec")
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
if msg.Channel != "system" {
|
||||
t.Errorf("expected channel=system, got %q", msg.Channel)
|
||||
if cbResult == nil {
|
||||
t.Fatal("callback was not invoked")
|
||||
}
|
||||
if msg.ChatID != "telegram:chat-123" {
|
||||
t.Errorf("expected chatID=telegram:chat-123, got %q", msg.ChatID)
|
||||
if !strings.Contains(cbResult.ForLLM, "bg_output") {
|
||||
t.Errorf("expected output in callback result: %s", cbResult.ForLLM)
|
||||
}
|
||||
if !strings.Contains(msg.Content, "bg_output") {
|
||||
t.Errorf("expected output in message content: %s", msg.Content)
|
||||
}
|
||||
if !strings.Contains(msg.Content, "completed") {
|
||||
t.Errorf("expected 'completed' in message content: %s", msg.Content)
|
||||
if !strings.Contains(cbResult.ForLLM, "completed") {
|
||||
t.Errorf("expected 'completed' in callback result: %s", cbResult.ForLLM)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecTool_BackgroundBlockedCommand(t *testing.T) {
|
||||
msgBus := bus.NewMessageBus()
|
||||
defer msgBus.Close()
|
||||
|
||||
tool, err := NewExecToolWithConfig(t.TempDir(), false, nil, msgBus)
|
||||
tool, err := NewExecTool(t.TempDir(), false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
ctx := WithToolContext(context.Background(), "telegram", "chat-123")
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(1)
|
||||
|
||||
cb := func(_ context.Context, _ *ToolResult) {}
|
||||
result := tool.ExecuteAsync(ctx, map[string]any{
|
||||
var cbResult *ToolResult
|
||||
cb := func(_ context.Context, result *ToolResult) {
|
||||
cbResult = result
|
||||
wg.Done()
|
||||
}
|
||||
|
||||
result := tool.ExecuteAsync(context.Background(), map[string]any{
|
||||
"command": "sudo rm -rf /",
|
||||
"background": true,
|
||||
}, cb)
|
||||
|
|
@ -115,13 +110,13 @@ func TestExecTool_BackgroundBlockedCommand(t *testing.T) {
|
|||
t.Fatal("expected async result even for blocked commands")
|
||||
}
|
||||
|
||||
msg, ok := msgBus.ConsumeInbound(ctx)
|
||||
if !ok {
|
||||
t.Fatal("expected inbound message from background exec")
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
if !strings.Contains(msg.Content, "failed") {
|
||||
t.Errorf("expected 'failed' in message content: %s", msg.Content)
|
||||
if cbResult == nil {
|
||||
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() {
|
||||
cfg := &config.Config{}
|
||||
cfg.Tools.Exec.EnableDenyPatterns = boolPtr(false)
|
||||
_, err := NewExecToolWithConfig(t.TempDir(), false, cfg, nil)
|
||||
_, err := NewExecToolWithConfig(t.TempDir(), false, cfg)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue