refactor(toolsdk): simplify workspace bash background execution with nohup

ThomasK33 · ThomasK33 · commit d6ae953512b8 · 2025-07-28T18:17:27.000+02:00
Change-Id: If67c30717158bdd84e9f733b56365af7c8d0b51a
Signed-off-by: Thomas Kosiewski &lt;tk@coder.com&gt;
diff --git a/codersdk/toolsdk/bash.go b/codersdk/toolsdk/bash.go
@@ -3,8 +3,6 @@ package toolsdk
 import (
 	"bytes"
 	"context"
-	_ "embed"
-	"encoding/base64"
 	"errors"
 	"fmt"
 	"io"
@@ -20,7 +18,6 @@ import (
 	"github.com/coder/coder/v2/cli/cliui"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/codersdk/workspacesdk"
-	"github.com/coder/coder/v2/cryptorand"
 )
 
 type WorkspaceBashArgs struct {
@@ -35,9 +32,6 @@ type WorkspaceBashResult struct {
 	ExitCode int    `json:"exit_code"`
 }
 
-//go:embed resources/background.sh
-var backgroundScript string
-
 var WorkspaceBash = Tool[WorkspaceBashArgs, WorkspaceBashResult]{
 	Tool: aisdk.Tool{
 		Name: ToolNameWorkspaceBash,
@@ -57,6 +51,9 @@ The workspace parameter supports various formats:
 The timeout_ms parameter specifies the command timeout in milliseconds (defaults to 60000ms, maximum of 300000ms).
 If the command times out, all output captured up to that point is returned with a cancellation message.
 
+For background commands (background: true), output is captured until the timeout is reached, then the command
+continues running in the background. The captured output is returned as the result.
+
 Examples:
 - workspace: "my-workspace", command: "ls -la"
 - workspace: "john/dev-env", command: "git status", timeout_ms: 30000
@@ -80,7 +77,7 @@ Examples:
 				},
 				"background": map[string]any{
 					"type":        "boolean",
-					"description": "Whether to run the command in the background. The command will not be affected by the timeout.",
+					"description": "Whether to run the command in the background. Output is captured until timeout, then the command continues running in the background.",
 				},
 			},
 			Required: []string{"workspace", "command"},
@@ -155,35 +152,28 @@ Examples:
 		}
 		command := args.Command
 		if args.Background {
-			// Background commands are not affected by the timeout
-			timeoutMs = defaultTimeoutMs
-			encodedCommand := base64.StdEncoding.EncodeToString([]byte(args.Command))
-			encodedScript := base64.StdEncoding.EncodeToString([]byte(backgroundScript))
-			commandID, err := cryptorand.StringCharset(cryptorand.Human, 8)
-			if err != nil {
-				return WorkspaceBashResult{}, xerrors.Errorf("failed to generate command ID: %w", err)
-			}
-			command = fmt.Sprintf(
-				"ARG_COMMAND=\"$(echo -n %s | base64 -d)\" ARG_COMMAND_ID=%s bash -c \"$(echo -n %s | base64 -d)\"",
-				encodedCommand,
-				commandID,
-				encodedScript,
-			)
+			// For background commands, use nohup to detach the process
+			// We redirect stdin to prevent hanging and combine stdout/stderr
+			command = fmt.Sprintf("nohup bash -c %q </dev/null 2>&1", args.Command)
 		}
 
-		// Create context with timeout
-		ctx, cancel = context.WithTimeout(ctx, time.Duration(timeoutMs)*time.Millisecond)
-		defer cancel()
+		// Create context with command timeout (replace the broader MCP timeout)
+		commandCtx, commandCancel := context.WithTimeout(ctx, time.Duration(timeoutMs)*time.Millisecond)
+		defer commandCancel()
 
 		// Execute command with timeout handling
-		output, err := executeCommandWithTimeout(ctx, session, command)
+		output, err := executeCommandWithTimeout(commandCtx, session, command)
 		outputStr := strings.TrimSpace(string(output))
 
 		// Handle command execution results
 		if err != nil {
 			// Check if the command timed out
-			if errors.Is(context.Cause(ctx), context.DeadlineExceeded) {
-				outputStr += "\nCommand canceled due to timeout"
+			if errors.Is(context.Cause(commandCtx), context.DeadlineExceeded) {
+				if args.Background {
+					outputStr += "\nCommand continues running in background"
+				} else {
+					outputStr += "\nCommand canceled due to timeout"
+				}
 				return WorkspaceBashResult{
 					Output:   outputStr,
 					ExitCode: 124,
@@ -417,21 +407,27 @@ func executeCommandWithTimeout(ctx context.Context, session *gossh.Session, comm
 		return safeWriter.Bytes(), err
 	case <-ctx.Done():
 		// Context was canceled (timeout or other cancellation)
-		// Close the session to stop the command
-		_ = session.Close()
+		// Close the session to stop the command, but handle errors gracefully
+		closeErr := session.Close()
 
-		// Give a brief moment to collect any remaining output
-		timer := time.NewTimer(50 * time.Millisecond)
+		// Give a brief moment to collect any remaining output and for goroutines to finish
+		timer := time.NewTimer(100 * time.Millisecond)
 		defer timer.Stop()
 
 		select {
 		case <-timer.C:
 			// Timer expired, return what we have
+			break
 		case err := <-done:
 			// Command finished during grace period
-			return safeWriter.Bytes(), err
+			if closeErr == nil {
+				return safeWriter.Bytes(), err
+			}
+			// If session close failed, prioritize the context error
+			break
 		}
 
+		// Return the collected output with the context error
 		return safeWriter.Bytes(), context.Cause(ctx)
 	}
 }
@@ -451,5 +447,9 @@ func (sw *syncWriter) Write(p []byte) (n int, err error) {
 func (sw *syncWriter) Bytes() []byte {
 	sw.mu.Lock()
 	defer sw.mu.Unlock()
-	return sw.w.Bytes()
+	// Return a copy to prevent race conditions with the underlying buffer
+	b := sw.w.Bytes()
+	result := make([]byte, len(b))
+	copy(result, b)
+	return result
 }
diff --git a/codersdk/toolsdk/bash_test.go b/codersdk/toolsdk/bash_test.go
@@ -314,15 +314,15 @@ func TestWorkspaceBashTimeoutIntegration(t *testing.T) {
 
 		deps, err := toolsdk.NewDeps(client)
 		require.NoError(t, err)
-		ctx := context.Background()
 
 		args := toolsdk.WorkspaceBashArgs{
 			Workspace: workspace.Name,
 			Command:   `echo "normal command"`, // Quick command that should complete normally
 			TimeoutMs: 5000,                    // 5 second timeout - plenty of time
 		}
 
-		result, err := toolsdk.WorkspaceBash.Handler(ctx, deps, args)
+		// Use testTool to register the tool as tested and satisfy coverage validation
+		result, err := testTool(t, toolsdk.WorkspaceBash, deps, args)
 
 		// Should not error
 		require.NoError(t, err)
@@ -343,7 +343,7 @@ func TestWorkspaceBashTimeoutIntegration(t *testing.T) {
 func TestWorkspaceBashBackgroundIntegration(t *testing.T) {
 	t.Parallel()
 
-	t.Run("BackgroundCommandReturnsImmediately", func(t *testing.T) {
+	t.Run("BackgroundCommandCapturesOutput", func(t *testing.T) {
 		t.Parallel()
 
 		client, workspace, agentToken := setupWorkspaceForAgent(t)
@@ -359,8 +359,9 @@ func TestWorkspaceBashBackgroundIntegration(t *testing.T) {
 
 		args := toolsdk.WorkspaceBashArgs{
 			Workspace:  workspace.Name,
-			Command:    `echo "started" && sleep 5 && echo "completed"`, // Command that would take 5+ seconds
-			Background: true,                                            // Run in background
+			Command:    `echo "started" && sleep 60 && echo "completed"`, // Command that would take 60+ seconds
+			Background: true,                                             // Run in background
+			TimeoutMs:  2000,                                             // 2 second timeout
 		}
 
 		result, err := toolsdk.WorkspaceBash.Handler(t.Context(), deps, args)
@@ -370,18 +371,17 @@ func TestWorkspaceBashBackgroundIntegration(t *testing.T) {
 
 		t.Logf("Background result: exitCode=%d, output=%q", result.ExitCode, result.Output)
 
-		// Should have exit code 0 (background start successful)
-		require.Equal(t, 0, result.ExitCode)
+		// Should have exit code 124 (timeout) since command times out
+		require.Equal(t, 124, result.ExitCode)
 
-		// Should contain PID and log path info, not the actual command output
-		require.Contains(t, result.Output, "Command started with PID:")
-		require.Contains(t, result.Output, "Log path: /tmp/mcp-bg/")
+		// Should capture output up to timeout point
+		require.Contains(t, result.Output, "started", "Should contain output captured before timeout")
+
+		// Should NOT contain the second echo (it never executed due to timeout)
+		require.NotContains(t, result.Output, "completed", "Should not contain output after timeout")
 
-		// Should NOT contain the actual command output since it runs in background
-		// The command was `echo "started" && sleep 5 && echo "completed"`
-		// So we check that the quoted strings don't appear in the output
-		require.NotContains(t, result.Output, `"started"`, "Should not contain command output in background mode")
-		require.NotContains(t, result.Output, `"completed"`, "Should not contain command output in background mode")
+		// Should contain background continuation message
+		require.Contains(t, result.Output, "Command continues running in background")
 	})
 
 	t.Run("BackgroundVsNormalExecution", func(t *testing.T) {
@@ -425,14 +425,12 @@ func TestWorkspaceBashBackgroundIntegration(t *testing.T) {
 		t.Logf("Normal result: %q", normalResult.Output)
 		t.Logf("Background result: %q", backgroundResult.Output)
 
-		// Background mode should return PID/log info, not the actual output
+		// Background mode should also return the actual output since command completes quickly
 		require.Equal(t, 0, backgroundResult.ExitCode)
-		require.Contains(t, backgroundResult.Output, "Command started with PID:")
-		require.Contains(t, backgroundResult.Output, "Log path: /tmp/mcp-bg/")
-		require.NotContains(t, backgroundResult.Output, "hello world")
+		require.Equal(t, "hello world", backgroundResult.Output)
 	})
 
-	t.Run("BackgroundIgnoresTimeout", func(t *testing.T) {
+	t.Run("BackgroundCommandContinuesAfterTimeout", func(t *testing.T) {
 		t.Parallel()
 
 		client, workspace, agentToken := setupWorkspaceForAgent(t)
@@ -448,36 +446,35 @@ func TestWorkspaceBashBackgroundIntegration(t *testing.T) {
 
 		args := toolsdk.WorkspaceBashArgs{
 			Workspace:  workspace.Name,
-			Command:    `sleep 1 && echo "done" > /tmp/done`, // Command that would normally timeout
-			TimeoutMs:  1,                                    // 1 ms timeout (shorter than command duration)
-			Background: true,                                 // But running in background should ignore timeout
+			Command:    `echo "started" && sleep 10 && echo "done" > /tmp/bg-test-done`, // Command that will timeout but continue
+			TimeoutMs:  5000,                                                            // 5000ms timeout (shorter than command duration)
+			Background: true,                                                            // Run in background
 		}
 
 		result, err := toolsdk.WorkspaceBash.Handler(t.Context(), deps, args)
 
-		// Should not error and should not timeout
+		// Should not error but should timeout
 		require.NoError(t, err)
 
 		t.Logf("Background with timeout result: exitCode=%d, output=%q", result.ExitCode, result.Output)
 
-		// Should have exit code 0 (background start successful)
-		require.Equal(t, 0, result.ExitCode)
+		// Should have timeout exit code
+		require.Equal(t, 124, result.ExitCode)
 
-		// Should return PID/log info, indicating the background command started successfully
-		require.Contains(t, result.Output, "Command started with PID:")
-		require.Contains(t, result.Output, "Log path: /tmp/mcp-bg/")
+		// Should capture output before timeout
+		require.Contains(t, result.Output, "started", "Should contain output captured before timeout")
 
-		// Should NOT contain timeout message since background mode ignores timeout
-		require.NotContains(t, result.Output, "Command canceled due to timeout")
+		// Should contain background continuation message
+		require.Contains(t, result.Output, "Command continues running in background")
 
-		// Wait for the background command to complete
+		// Wait for the background command to complete (even though SSH session timed out)
 		require.Eventually(t, func() bool {
-			args := toolsdk.WorkspaceBashArgs{
+			checkArgs := toolsdk.WorkspaceBashArgs{
 				Workspace: workspace.Name,
-				Command:   `cat /tmp/done`,
+				Command:   `cat /tmp/bg-test-done 2>/dev/null || echo "not found"`,
 			}
-			result, err := toolsdk.WorkspaceBash.Handler(t.Context(), deps, args)
-			return err == nil && result.Output == "done"
-		}, testutil.WaitMedium, testutil.IntervalMedium)
+			checkResult, err := toolsdk.WorkspaceBash.Handler(t.Context(), deps, checkArgs)
+			return err == nil && checkResult.Output == "done"
+		}, testutil.WaitMedium, testutil.IntervalMedium, "Background command should continue running and complete after timeout")
 	})
 }
diff --git a/codersdk/toolsdk/resources/background.sh b/codersdk/toolsdk/resources/background.sh
