chore: small fixes & killall on shutdown

kujtimiihoxha created

Change summary

internal/agent/tools/bash_kill.go   | 26 ++----------
internal/agent/tools/bash_kill.md   | 18 +++++++++
internal/agent/tools/bash_output.go | 27 ++-----------
internal/agent/tools/bash_output.md | 19 +++++++++
internal/app/app.go                 |  4 ++
internal/shell/background.go        | 16 ++++++++
internal/shell/background_test.go   | 62 +++++++++++++++++++++++++++++++
7 files changed, 129 insertions(+), 43 deletions(-)

Detailed changes

internal/agent/tools/bash_kill.go 🔗

@@ -2,6 +2,7 @@ package tools
 
 import (
 	"context"
+	_ "embed"
 	"fmt"
 
 	"charm.land/fantasy"
@@ -12,6 +13,9 @@ const (
 	BashKillToolName = "bash_kill"
 )
 
+//go:embed bash_kill.md
+var bashKillDescription []byte
+
 type BashKillParams struct {
 	ShellID string `json:"shell_id" description:"The ID of the background shell to terminate"`
 }
@@ -20,30 +24,10 @@ type BashKillResponseMetadata struct {
 	ShellID string `json:"shell_id"`
 }
 
-const bashKillDescription = `Terminates a background shell process.
-
-<usage>
-- Provide the shell ID returned from a background bash execution
-- Cancels the running process and cleans up resources
-</usage>
-
-<features>
-- Stop long-running background processes
-- Clean up completed background shells
-- Immediately terminates the process
-</features>
-
-<tips>
-- Use this when you need to stop a background process
-- The process is terminated immediately (similar to SIGTERM)
-- After killing, the shell ID becomes invalid
-</tips>
-`
-
 func NewBashKillTool() fantasy.AgentTool {
 	return fantasy.NewAgentTool(
 		BashKillToolName,
-		bashKillDescription,
+		string(bashKillDescription),
 		func(ctx context.Context, params BashKillParams, call fantasy.ToolCall) (fantasy.ToolResponse, error) {
 			if params.ShellID == "" {
 				return fantasy.NewTextErrorResponse("missing shell_id"), nil

internal/agent/tools/bash_kill.md 🔗

@@ -0,0 +1,18 @@
+Terminates a background shell process.
+
+<usage>
+- Provide the shell ID returned from a background bash execution
+- Cancels the running process and cleans up resources
+</usage>
+
+<features>
+- Stop long-running background processes
+- Clean up completed background shells
+- Immediately terminates the process
+</features>
+
+<tips>
+- Use this when you need to stop a background process
+- The process is terminated immediately (similar to SIGTERM)
+- After killing, the shell ID becomes invalid
+</tips>

internal/agent/tools/bash_output.go 🔗

@@ -2,6 +2,7 @@ package tools
 
 import (
 	"context"
+	_ "embed"
 	"fmt"
 	"strings"
 
@@ -13,6 +14,9 @@ const (
 	BashOutputToolName = "bash_output"
 )
 
+//go:embed bash_output.md
+var bashOutputDescription []byte
+
 type BashOutputParams struct {
 	ShellID string `json:"shell_id" description:"The ID of the background shell to retrieve output from"`
 }
@@ -23,31 +27,10 @@ type BashOutputResponseMetadata struct {
 	WorkingDirectory string `json:"working_directory"`
 }
 
-const bashOutputDescription = `Retrieves the current output from a background shell.
-
-<usage>
-- Provide the shell ID returned from a background bash execution
-- Returns the current stdout and stderr output
-- Indicates whether the shell has completed execution
-</usage>
-
-<features>
-- View output from running background processes
-- Check if background process has completed
-- Get cumulative output from process start
-</features>
-
-<tips>
-- Use this to monitor long-running processes
-- Check the 'done' status to see if process completed
-- Can be called multiple times to view incremental output
-</tips>
-`
-
 func NewBashOutputTool() fantasy.AgentTool {
 	return fantasy.NewAgentTool(
 		BashOutputToolName,
-		bashOutputDescription,
+		string(bashOutputDescription),
 		func(ctx context.Context, params BashOutputParams, call fantasy.ToolCall) (fantasy.ToolResponse, error) {
 			if params.ShellID == "" {
 				return fantasy.NewTextErrorResponse("missing shell_id"), nil

internal/agent/tools/bash_output.md 🔗

@@ -0,0 +1,19 @@
+Retrieves the current output from a background shell.
+
+<usage>
+- Provide the shell ID returned from a background bash execution
+- Returns the current stdout and stderr output
+- Indicates whether the shell has completed execution
+</usage>
+
+<features>
+- View output from running background processes
+- Check if background process has completed
+- Get cumulative output from process start
+</features>
+
+<tips>
+- Use this to monitor long-running processes
+- Check the 'done' status to see if process completed
+- Can be called multiple times to view incremental output
+</tips>

internal/app/app.go 🔗

@@ -24,6 +24,7 @@ import (
 	"github.com/charmbracelet/crush/internal/permission"
 	"github.com/charmbracelet/crush/internal/pubsub"
 	"github.com/charmbracelet/crush/internal/session"
+	"github.com/charmbracelet/crush/internal/shell"
 	"github.com/charmbracelet/x/ansi"
 )
 
@@ -325,6 +326,9 @@ func (app *App) Shutdown() {
 		app.AgentCoordinator.CancelAll()
 	}
 
+	// Kill all background shells.
+	shell.GetBackgroundShellManager().KillAll()
+
 	// Shutdown all LSP clients.
 	for name, client := range app.LSPClients.Seq2() {
 		shutdownCtx, cancel := context.WithTimeout(app.globalCtx, 5*time.Second)

internal/shell/background.go 🔗

@@ -121,6 +121,22 @@ func (m *BackgroundShellManager) List() []string {
 	return ids
 }
 
+// KillAll terminates all background shells.
+func (m *BackgroundShellManager) KillAll() {
+	m.mu.Lock()
+	shells := make([]*BackgroundShell, 0, len(m.shells))
+	for _, shell := range m.shells {
+		shells = append(shells, shell)
+	}
+	m.shells = make(map[string]*BackgroundShell)
+	m.mu.Unlock()
+
+	for _, shell := range shells {
+		shell.cancel()
+		<-shell.done
+	}
+}
+
 // GetOutput returns the current output of a background shell.
 func (bs *BackgroundShell) GetOutput() (stdout string, stderr string, done bool, err error) {
 	bs.mu.RLock()

internal/shell/background_test.go 🔗

@@ -212,3 +212,65 @@ func TestBackgroundShellManager_List(t *testing.T) {
 	manager.Kill(bgShell1.ID)
 	manager.Kill(bgShell2.ID)
 }
+
+func TestBackgroundShellManager_KillAll(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	workingDir := t.TempDir()
+	manager := GetBackgroundShellManager()
+
+	// Start multiple long-running shells
+	shell1, err := manager.Start(ctx, workingDir, nil, "sleep 10")
+	if err != nil {
+		t.Fatalf("failed to start shell 1: %v", err)
+	}
+
+	shell2, err := manager.Start(ctx, workingDir, nil, "sleep 10")
+	if err != nil {
+		t.Fatalf("failed to start shell 2: %v", err)
+	}
+
+	shell3, err := manager.Start(ctx, workingDir, nil, "sleep 10")
+	if err != nil {
+		t.Fatalf("failed to start shell 3: %v", err)
+	}
+
+	// Verify shells are running
+	if shell1.IsDone() || shell2.IsDone() || shell3.IsDone() {
+		t.Error("shells should not be done yet")
+	}
+
+	// Kill all shells
+	manager.KillAll()
+
+	// Verify all shells are done
+	if !shell1.IsDone() {
+		t.Error("shell1 should be done after KillAll")
+	}
+	if !shell2.IsDone() {
+		t.Error("shell2 should be done after KillAll")
+	}
+	if !shell3.IsDone() {
+		t.Error("shell3 should be done after KillAll")
+	}
+
+	// Verify they're removed from the manager
+	if _, ok := manager.Get(shell1.ID); ok {
+		t.Error("shell1 should be removed from manager")
+	}
+	if _, ok := manager.Get(shell2.ID); ok {
+		t.Error("shell2 should be removed from manager")
+	}
+	if _, ok := manager.Get(shell3.ID); ok {
+		t.Error("shell3 should be removed from manager")
+	}
+
+	// Verify list is empty (or doesn't contain our shells)
+	ids := manager.List()
+	for _, id := range ids {
+		if id == shell1.ID || id == shell2.ID || id == shell3.ID {
+			t.Errorf("shell %s should not be in list after KillAll", id)
+		}
+	}
+}