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:
Boris Bliznioukov 2026-03-05 13:29:09 +01:00
parent bf1b07d407
commit 970a56801b
No known key found for this signature in database
7 changed files with 81 additions and 83 deletions

View file

@ -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 {

View file

@ -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)
} }

View file

@ -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

View file

@ -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)
} }

View file

@ -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") {

View file

@ -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)))
} }

View file

@ -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)
} }