claudetool/bash: include partial output in error message for timeouts

Co-Authored-By: sketch <hello@sketch.dev>
Change-ID: se1cb20e743a878d4k
diff --git a/claudetool/bash.go b/claudetool/bash.go
index adbd12c..827235a 100644
--- a/claudetool/bash.go
+++ b/claudetool/bash.go
@@ -138,12 +138,7 @@
 
 	// For foreground commands, use executeBash
 	out, execErr := executeBash(ctx, req)
-	// If there's a timeout error, we still want to return the partial output
 	if execErr != nil {
-		if out != "" && strings.Contains(execErr.Error(), "timed out") {
-			// Return both the partial output and the error
-			return llm.TextContent(out), execErr
-		}
 		return nil, execErr
 	}
 	return llm.TextContent(out), nil
@@ -186,15 +181,6 @@
 	err := cmd.Wait()
 	close(done)
 
-	if execCtx.Err() == context.DeadlineExceeded {
-		// Get the partial output that was captured before the timeout
-		partialOutput := output.String()
-		// Truncate if the output is too large
-		if len(partialOutput) > maxBashOutputLength {
-			partialOutput = partialOutput[:maxBashOutputLength] + "\n[output truncated due to size]\n"
-		}
-		return partialOutput, fmt.Errorf("command timed out after %s - partial output included", req.timeout())
-	}
 	longOutput := output.Len() > maxBashOutputLength
 	var outstr string
 	if longOutput {
@@ -206,6 +192,15 @@
 		outstr = output.String()
 	}
 
+	if execCtx.Err() == context.DeadlineExceeded {
+		// Get the partial output that was captured before the timeout
+		partialOutput := output.String()
+		// Truncate if the output is too large
+		if len(partialOutput) > maxBashOutputLength {
+			partialOutput = partialOutput[:maxBashOutputLength] + "\n[output truncated due to size]\n"
+		}
+		return "", fmt.Errorf("command timed out after %s\nCommand output (until it timed out):\n%s", req.timeout(), outstr)
+	}
 	if err != nil {
 		return "", fmt.Errorf("command failed: %w\n%s", err, outstr)
 	}
diff --git a/test/timeout_test.go b/test/timeout_test.go
index 87ab10a..6e78098 100644
--- a/test/timeout_test.go
+++ b/test/timeout_test.go
@@ -6,7 +6,6 @@
 	"testing"
 
 	"sketch.dev/claudetool"
-	"sketch.dev/llm"
 )
 
 func TestBashTimeout(t *testing.T) {
@@ -42,30 +41,25 @@
 		t.Errorf("Error doesn't mention timeout: %v", err)
 	}
 
-	// Check that we got partial output despite the error
-	if len(result) == 0 {
-		t.Fatalf("Expected partial output, got empty result")
+	// No output should be returned directly, it should be in the error message
+	if len(result) > 0 {
+		t.Fatalf("Expected no direct output, got: %v", result)
 	}
 
-	// Verify the error mentions that partial output is included
-	if !containsString(err.Error(), "partial output included") {
-		t.Errorf("Error should mention that partial output is included: %v", err)
+	// The error should contain the partial output
+	errorMsg := err.Error()
+	if !containsString(errorMsg, "Starting command") || !containsString(errorMsg, "should appear in partial output") {
+		t.Errorf("Error should contain the partial output: %v", errorMsg)
 	}
 
-	// The partial output should contain the initial output but not the text after sleep
-	text := ""
-	for _, content := range result {
-		if content.Type == llm.ContentTypeText {
-			text += content.Text
-		}
+	// The error should indicate a timeout
+	if !containsString(errorMsg, "timed out") {
+		t.Errorf("Error should indicate a timeout: %v", errorMsg)
 	}
 
-	if !containsString(text, "Starting command") || !containsString(text, "should appear in partial output") {
-		t.Errorf("Partial output is missing expected content: %s", text)
-	}
-
-	if containsString(text, "shouldn't appear") {
-		t.Errorf("Partial output contains unexpected content (after timeout): %s", text)
+	// The error should not contain the output that would appear after the sleep
+	if containsString(err.Error(), "shouldn't appear") {
+		t.Errorf("Error contains output that should not have been captured (after timeout): %s", err.Error())
 	}
 }