Terminate process group

pkg/tools/shell.go

@@ -1,14 +1,9 @@
 package tools
 
 import (
-	"bytes"
-	"context"
 	"fmt"
-	"os"
-	"os/exec"
 	"path/filepath"
 	"regexp"
-	"runtime"
 	"strings"
 	"time"
 )
@@ -67,87 +62,6 @@ func (t *ExecTool) Parameters() map[string]interface{} {
 	}
 }
 
-func (t *ExecTool) Execute(ctx context.Context, args map[string]interface{}) *ToolResult {
-	command, ok := args["command"].(string)
-	if !ok {
-		return ErrorResult("command is required")
-	}
-
-	cwd := t.workingDir
-	if wd, ok := args["working_dir"].(string); ok && wd != "" {
-		cwd = wd
-	}
-
-	if cwd == "" {
-		wd, err := os.Getwd()
-		if err == nil {
-			cwd = wd
-		}
-	}
-
-	if guardError := t.guardCommand(command, cwd); guardError != "" {
-		return ErrorResult(guardError)
-	}
-
-	cmdCtx, cancel := context.WithTimeout(ctx, t.timeout)
-	defer cancel()
-
-	var cmd *exec.Cmd
-	if runtime.GOOS == "windows" {
-		cmd = exec.CommandContext(cmdCtx, "powershell", "-NoProfile", "-NonInteractive", "-Command", command)
-	} else {
-		cmd = exec.CommandContext(cmdCtx, "sh", "-c", command)
-	}
-	if cwd != "" {
-		cmd.Dir = cwd
-	}
-
-	var stdout, stderr bytes.Buffer
-	cmd.Stdout = &stdout
-	cmd.Stderr = &stderr
-
-	err := cmd.Run()
-	output := stdout.String()
-	if stderr.Len() > 0 {
-		output += "\nSTDERR:\n" + stderr.String()
-	}
-
-	if err != nil {
-		if cmdCtx.Err() == context.DeadlineExceeded {
-			msg := fmt.Sprintf("Command timed out after %v", t.timeout)
-			return &ToolResult{
-				ForLLM:  msg,
-				ForUser: msg,
-				IsError: true,
-			}
-		}
-		output += fmt.Sprintf("\nExit code: %v", err)
-	}
-
-	if output == "" {
-		output = "(no output)"
-	}
-
-	maxLen := 10000
-	if len(output) > maxLen {
-		output = output[:maxLen] + fmt.Sprintf("\n... (truncated, %d more chars)", len(output)-maxLen)
-	}
-
-	if err != nil {
-		return &ToolResult{
-			ForLLM:  output,
-			ForUser: output,
-			IsError: true,
-		}
-	}
-
-	return &ToolResult{
-		ForLLM:  output,
-		ForUser: output,
-		IsError: false,
-	}
-}
-
 func (t *ExecTool) guardCommand(command, cwd string) string {
 	cmd := strings.TrimSpace(command)
 	lower := strings.ToLower(cmd)

pkg/tools/shell_unix.go

@@ -0,0 +1,121 @@
+//go:build !windows
+
+package tools
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"os"
+	"os/exec"
+	"syscall"
+)
+
+func (t *ExecTool) Execute(ctx context.Context, args map[string]interface{}) *ToolResult {
+	command, ok := args["command"].(string)
+	if !ok {
+		return ErrorResult("command is required")
+	}
+
+	cwd := t.workingDir
+	if wd, ok := args["working_dir"].(string); ok && wd != "" {
+		cwd = wd
+	}
+
+	if cwd == "" {
+		wd, err := os.Getwd()
+		if err == nil {
+			cwd = wd
+		}
+	}
+
+	if guardError := t.guardCommand(command, cwd); guardError != "" {
+		return ErrorResult(guardError)
+	}
+
+	cmdCtx, cancel := context.WithTimeout(ctx, t.timeout)
+	defer cancel()
+
+	cmd := exec.CommandContext(cmdCtx, "sh", "-c", command)
+
+	cmd.SysProcAttr = &syscall.SysProcAttr{
+		Setpgid: true,
+	}
+
+	if cwd != "" {
+		cmd.Dir = cwd
+	}
+
+	var stdout, stderr bytes.Buffer
+	cmd.Stdout = &stdout
+	cmd.Stderr = &stderr
+
+	err := cmd.Start()
+	if err != nil {
+		output := stdout.String()
+		if stderr.Len() > 0 {
+			output += "\nSTDERR:\n" + stderr.String()
+		}
+		return &ToolResult{
+			ForLLM:  fmt.Sprintf("Failed to start command: %v\n%s", err, output),
+			ForUser: fmt.Sprintf("Failed to start command: %v\n%s", err, output),
+			IsError: true,
+		}
+	}
+
+	done := make(chan error, 1)
+	go func() {
+		done <- cmd.Wait()
+	}()
+
+	select {
+	case <-cmdCtx.Done():
+		if cmd.Process != nil {
+			pgid, err := syscall.Getpgid(cmd.Process.Pid)
+			if err == nil {
+				syscall.Kill(-pgid, syscall.SIGKILL)
+			}
+		}
+		err = cmdCtx.Err()
+	case err = <-done:
+	}
+	output := stdout.String()
+	if stderr.Len() > 0 {
+		output += "\nSTDERR:\n" + stderr.String()
+	}
+
+	if err != nil {
+		if cmdCtx.Err() == context.DeadlineExceeded {
+			msg := fmt.Sprintf("Command timed out after %v", t.timeout)
+			return &ToolResult{
+				ForLLM:  msg,
+				ForUser: msg,
+				IsError: true,
+			}
+		}
+		output += fmt.Sprintf("\nExit code: %v", err)
+	}
+
+	if output == "" {
+		output = "(no output)"
+	}
+
+	maxLen := 10000
+	if len(output) > maxLen {
+		output = output[:maxLen] + fmt.Sprintf("\n... (truncated, %d more chars)", len(output)-maxLen)
+	}
+
+	if err != nil {
+		return &ToolResult{
+			ForLLM:  output,
+			ForUser: output,
+			IsError: true,
+		}
+	}
+
+	return &ToolResult{
+		ForLLM:  output,
+		ForUser: output,
+		IsError: false,
+	}
+}

pkg/tools/shell_windows.go

@@ -0,0 +1,113 @@
+//go:build windows
+
+package tools
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"os"
+	"os/exec"
+)
+
+func (t *ExecTool) Execute(ctx context.Context, args map[string]interface{}) *ToolResult {
+	command, ok := args["command"].(string)
+	if !ok {
+		return ErrorResult("command is required")
+	}
+
+	cwd := t.workingDir
+	if wd, ok := args["working_dir"].(string); ok && wd != "" {
+		cwd = wd
+	}
+
+	if cwd == "" {
+		wd, err := os.Getwd()
+		if err == nil {
+			cwd = wd
+		}
+	}
+
+	if guardError := t.guardCommand(command, cwd); guardError != "" {
+		return ErrorResult(guardError)
+	}
+
+	cmdCtx, cancel := context.WithTimeout(ctx, t.timeout)
+	defer cancel()
+
+	cmd := exec.CommandContext(cmdCtx, "powershell", "-NoProfile", "-NonInteractive", "-Command", command)
+
+	if cwd != "" {
+		cmd.Dir = cwd
+	}
+
+	var stdout, stderr bytes.Buffer
+	cmd.Stdout = &stdout
+	cmd.Stderr = &stderr
+
+	err := cmd.Start()
+	if err != nil {
+		output := stdout.String()
+		if stderr.Len() > 0 {
+			output += "\nSTDERR:\n" + stderr.String()
+		}
+		return &ToolResult{
+			ForLLM:  fmt.Sprintf("Failed to start command: %v\n%s", err, output),
+			ForUser: fmt.Sprintf("Failed to start command: %v\n%s", err, output),
+			IsError: true,
+		}
+	}
+
+	done := make(chan error, 1)
+	go func() {
+		done <- cmd.Wait()
+	}()
+
+	select {
+	case <-cmdCtx.Done():
+		if cmd.Process != nil {
+			cmd.Process.Kill()
+		}
+		err = cmdCtx.Err()
+	case err = <-done:
+	}
+	output := stdout.String()
+	if stderr.Len() > 0 {
+		output += "\nSTDERR:\n" + stderr.String()
+	}
+
+	if err != nil {
+		if cmdCtx.Err() == context.DeadlineExceeded {
+			msg := fmt.Sprintf("Command timed out after %v", t.timeout)
+			return &ToolResult{
+				ForLLM:  msg,
+				ForUser: msg,
+				IsError: true,
+			}
+		}
+		output += fmt.Sprintf("\nExit code: %v", err)
+	}
+
+	if output == "" {
+		output = "(no output)"
+	}
+
+	maxLen := 10000
+	if len(output) > maxLen {
+		output = output[:maxLen] + fmt.Sprintf("\n... (truncated, %d more chars)", len(output)-maxLen)
+	}
+
+	if err != nil {
+		return &ToolResult{
+			ForLLM:  output,
+			ForUser: output,
+			IsError: true,
+		}
+	}
+
+	return &ToolResult{
+		ForLLM:  output,
+		ForUser: output,
+		IsError: false,
+	}
+}