feat: combine default and user set banned list based on toggle

tauraamui created 4 months ago

Change summary

internal/agent/common_test.go |  2 +-
internal/agent/tools/bash.go  | 28 ++++++++++++++++++++++------
2 files changed, 23 insertions(+), 7 deletions(-)

Detailed changes

internal/agent/common_test.go

@@ -190,7 +190,7 @@ func coderAgent(r *vcr.Recorder, env fakeEnv, large, small fantasy.LanguageModel
 	}
 
 	allTools := []fantasy.AgentTool{
-		tools.NewBashTool(env.permissions, env.workingDir, cfg.Options.Attribution, modelName),
+		tools.NewBashTool(env.permissions, env.workingDir, cfg.Options.Attribution, modelName, cfg.Tools.Bash),
 		tools.NewDownloadTool(env.permissions, env.workingDir, r.GetDefaultClient()),
 		tools.NewEditTool(env.lspClients, env.permissions, env.history, env.workingDir),
 		tools.NewMultiEditTool(env.lspClients, env.permissions, env.history, env.workingDir),

internal/agent/tools/bash.go

@@ -66,7 +66,7 @@ type bashDescriptionData struct {
 	ModelName       string
 }
 
-var bannedCommands = []string{
+var defaultBannedCommands = []string{
 	// Network/Download tools
 	"alias",
 	"aria2c",
@@ -140,7 +140,7 @@ var bannedCommands = []string{
 }
 
 func bashDescription(attribution *config.Attribution, modelName string) string {
-	bannedCommandsStr := strings.Join(bannedCommands, ", ")
+	bannedCommandsStr := strings.Join(defaultBannedCommands, ", ")
 	var out bytes.Buffer
 	if err := bashDescriptionTpl.Execute(&out, bashDescriptionData{
 		BannedCommands:  bannedCommandsStr,
@@ -154,7 +154,7 @@ func bashDescription(attribution *config.Attribution, modelName string) string {
 	return out.String()
 }
 
-func blockFuncs() []shell.BlockFunc {
+func blockFuncs(bannedCommands []string) []shell.BlockFunc {
 	return []shell.BlockFunc{
 		shell.CommandsBlocker(bannedCommands),
 
@@ -186,7 +186,23 @@ func blockFuncs() []shell.BlockFunc {
 	}
 }
 
-func NewBashTool(permissions permission.Service, workingDir string, attribution *config.Attribution, modelName string) fantasy.AgentTool {
+func resolveBlockFuncs(cfg config.ToolBash) []shell.BlockFunc {
+	bannedCommands := cfg.BannedCommands
+	if !cfg.DisableDefaults {
+		if len(bannedCommands) == 0 {
+			return blockFuncs(defaultBannedCommands)
+		}
+		bannedCommands = append(bannedCommands, defaultBannedCommands...)
+	}
+	return blockFuncs(bannedCommands)
+}
+
+func NewBashTool(
+	permissions permission.Service,
+	workingDir string, attribution *config.Attribution,
+	modelName string,
+	bashConfig config.ToolBash,
+) fantasy.AgentTool {
 	return fantasy.NewAgentTool(
 		BashToolName,
 		string(bashDescription(attribution, modelName)),
@@ -237,7 +253,7 @@ func NewBashTool(permissions permission.Service, workingDir string, attribution
 				bgManager := shell.GetBackgroundShellManager()
 				bgManager.Cleanup()
 				// Use background context so it continues after tool returns
-				bgShell, err := bgManager.Start(context.Background(), execWorkingDir, blockFuncs(), params.Command, params.Description)
+				bgShell, err := bgManager.Start(context.Background(), execWorkingDir, resolveBlockFuncs(bashConfig), params.Command, params.Description)
 				if err != nil {
 					return fantasy.ToolResponse{}, fmt.Errorf("error starting background shell: %w", err)
 				}
@@ -292,7 +308,7 @@ func NewBashTool(permissions permission.Service, workingDir string, attribution
 			// Start with detached context so it can survive if moved to background
 			bgManager := shell.GetBackgroundShellManager()
 			bgManager.Cleanup()
-			bgShell, err := bgManager.Start(context.Background(), execWorkingDir, blockFuncs(), params.Command, params.Description)
+			bgShell, err := bgManager.Start(context.Background(), execWorkingDir, blockFuncs(defaultBannedCommands), params.Command, params.Description)
 			if err != nil {
 				return fantasy.ToolResponse{}, fmt.Errorf("error starting shell: %w", err)
 			}