diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 053643cbb..7452be8d1 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -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 { diff --git a/pkg/tools/cron.go b/pkg/tools/cron.go index 2ec44748d..6af0aa9e1 100644 --- a/pkg/tools/cron.go +++ b/pkg/tools/cron.go @@ -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) } diff --git a/pkg/tools/shell/risk.go b/pkg/tools/shell/risk.go index 7b6ab130c..0d1b2f376 100644 --- a/pkg/tools/shell/risk.go +++ b/pkg/tools/shell/risk.go @@ -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 diff --git a/pkg/tools/shell/runner.go b/pkg/tools/shell/runner.go index 85138e73f..6c2378fc2 100644 --- a/pkg/tools/shell/runner.go +++ b/pkg/tools/shell/runner.go @@ -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) } diff --git a/pkg/tools/shell/runner_windows_test.go b/pkg/tools/shell/runner_windows_test.go index 12c464277..4852d0de9 100644 --- a/pkg/tools/shell/runner_windows_test.go +++ b/pkg/tools/shell/runner_windows_test.go @@ -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") { diff --git a/pkg/tools/shell_tool.go b/pkg/tools/shell_tool.go index 3aad99006..03bcd1e2b 100644 --- a/pkg/tools/shell_tool.go +++ b/pkg/tools/shell_tool.go @@ -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))) } diff --git a/pkg/tools/shell_tool_test.go b/pkg/tools/shell_tool_test.go index 3b3c85a09..8b743f9d6 100644 --- a/pkg/tools/shell_tool_test.go +++ b/pkg/tools/shell_tool_test.go @@ -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) }