fix: subagent hangs — bash timeouts, process-group kill, watchdogs

cmd/late/main.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"context"
+	"errors"
 	"flag"
 	"fmt"
 	"late/internal/agent"
@@ -74,17 +75,23 @@ func (p pluginInlineTool) CallString(args json.RawMessage) string {
 	return fmt.Sprintf("Calling plugin tool %q...", p.name)
 }
 
+// subagentTimeoutUsage is single-sourced with the -subagent-timeout flag
+// registration in main() so tests can assert on the rendered help output.
+const subagentTimeoutUsage = "Max wall-clock time for one subagent run (0 = unlimited)"
+
 func main() {
 	// Parse flags
 	helpReq := flag.Bool("help", false, "Show help")
 	systemPromptReq := flag.String("system-prompt", "", "Set the system prompt (literal string)")
 	systemPromptFileReq := flag.String("system-prompt-file", "", "Set the system prompt from a file")
 	useToolsReq := flag.Bool("use-tools", true, "Enable tool usage (allows LLM to call tools)")
 	enableBashReq := flag.Bool("enable-bash", true, "Enable bash tool execution")
+	bashTimeout := flag.Duration("bash-timeout", 10*time.Minute, "Max wall-clock time for one bash tool call (0 = unlimited)")
 	injectCWDReq := flag.Bool("inject-cwd", true, "Replace ${{CWD}} in system prompt with current working directory")
 	enableSubagentsReq := flag.Bool("enable-subagents", true, "Enable subagent usage")
 	gemmaThinkingReq := flag.Bool("gemma-thinking", false, "Prepend <|think|> token to system prompt for Gemma 4 models")
 	subagentMaxTurns := flag.Int("subagent-max-turns", 500, "Maximum number of turns for subagents (default: 500)")
+	subagentTimeout := flag.Duration("subagent-timeout", appconfig.DefaultSubagentTimeout, subagentTimeoutUsage)
 	saveSubagentHistoriesReq := flag.Bool("save-subagent-histories", false, "Persist subagent conversation histories to disk (default: off)")
 	enableSqzReq := flag.Bool("enable-sqz", false, "Enable sqz context compression (if available)")
 	appendSystemPromptReq := flag.String("append-system-prompt", "", "Append text to the system prompt after processing")
@@ -126,6 +133,9 @@ func main() {
 	flag.Parse()
 
 	tool.SetSqzEnabled(*enableSqzReq)
+	// Shell tool bound: 0/negative (--bash-timeout=0) disables it —
+	// ShellTool.Execute treats a non-positive timeout as unbounded.
+	tool.SetShellTimeout(*bashTimeout)
 
 	if *versionReq {
 		fmt.Printf("late %s\n", common.Version)
@@ -375,6 +385,13 @@ func main() {
 	}
 	saveSubagentHistories := appconfig.ResolveSaveSubagentHistories(appConfig, saveSubagentHistoriesCLI, *saveSubagentHistoriesReq, storedSubagentHistoryPreference)
 
+	// Resolve the per-subagent wall-clock budget
+	// (explicit CLI flag > config.json subagent-timeout entry > default).
+	resolvedSubagentTimeout, subagentTimeoutWarning := resolveSubagentTimeBudget(flag.CommandLine, appConfig, *subagentTimeout)
+	if subagentTimeoutWarning != "" {
+		fmt.Fprintf(os.Stderr, "Warning: %s\n", subagentTimeoutWarning)
+	}
+
 	// Initialize Core Components
 	resolvedOpenAIConfig := appconfig.ResolveOpenAISettings(appConfig)
 	resolvedClientConfig := client.Config{
@@ -703,9 +720,47 @@ func main() {
 			}
 			child.SetMiddlewares(buildMiddlewares(pluginManager, p, child.Registry()))
 
+			// Per-subagent wall-clock budget layered on top of the parent
+			// context (agent.NewSubagentOrchestrator already wires the child
+			// to the parent's context). 0 disables the budget, so no wrapper
+			// is installed then: context.WithTimeout with a non-positive
+			// duration would expire immediately and kill the run.
+			runCtx := ctx
+			runCancel := func() {}
+			if resolvedSubagentTimeout > 0 {
+				runCtx, runCancel = context.WithTimeout(ctx, resolvedSubagentTimeout)
+			}
+			defer runCancel()
+			// SetContext is not part of common.Orchestrator; the factory
+			// always returns *orchestrator.BaseOrchestrator. If the concrete
+			// type ever changes, the child simply keeps the parent context
+			// and runs without a budget rather than failing the spawn.
+			if base, ok := child.(*orchestrator.BaseOrchestrator); ok {
+				base.SetContext(runCtx)
+			}
+
 			res, err := child.Execute("")
-			if err != nil {
-				return "", err
+			// Classify the termination BEFORE looking at err: the budget, the
+			// user's kill and crashes all surface differently here.
+			var cause string
+			switch {
+			case resolvedSubagentTimeout > 0 && runCtx.Err() == context.DeadlineExceeded:
+				cause = fmt.Sprintf("time budget exhausted (%s)", resolvedSubagentTimeout)
+			case errors.Is(err, context.Canceled) || child.IsStopRequested():
+				cause = "cancelled or killed by the user"
+			case err != nil:
+				cause = fmt.Sprintf("crashed: %v", err)
+			}
+			if cause != "" {
+				// Abnormal termination: hand the parent a pruned transcript so it
+				// can understand the cause and resume without redoing the work.
+				transcriptPath, terr := writeSubagentTranscript(child, agentType, goal, cause)
+				summary := lastActionPreview(child.History(), 500)
+				result := fmt.Sprintf("The %s subagent terminated abnormally (%s).\nFull pruned transcript: %s\nLast actions:\n%s", agentType, cause, transcriptPath, summary)
+				if terr != nil {
+					result += fmt.Sprintf("\n(transcript unavailable: %v)", terr)
+				}
+				return result, nil
 			}
 
 			if child.IsStopRequested() {
@@ -738,6 +793,22 @@ func deriveEffectiveSessionID(historyPath string) string {
 	}
 	return id
 }
+
+// resolveSubagentTimeBudget resolves the effective per-subagent wall-clock
+// budget: an explicitly-passed -subagent-timeout flag wins over the
+// config.json "subagent-timeout" entry, which wins over the
+// appconfig.DefaultSubagentTimeout default. Explicitness is detected on the
+// given FlagSet the same way main() detects --save-subagent-histories.
+// A zero or negative budget means unlimited.
+func resolveSubagentTimeBudget(fs *flag.FlagSet, cfg *appconfig.Config, flagValue time.Duration) (time.Duration, string) {
+	explicit := false
+	fs.Visit(func(f *flag.Flag) {
+		if f.Name == "subagent-timeout" {
+			explicit = true
+		}
+	})
+	return appconfig.ResolveSubagentTimeout(cfg, explicit, flagValue)
+}
 func newModelClient(ctx context.Context, setting appconfig.ModelSetting, enableImages bool, logitBias map[string]int) *client.Client {
 	c := client.NewClient(client.Config{
 		BaseURL:      setting.URL,

cmd/late/main_test.go

@@ -2,10 +2,12 @@ package main
 
 import (
 	"encoding/json"
+	"flag"
 	"net/http"
 	"net/http/httptest"
 	"os"
 	"path/filepath"
+	"strings"
 	"testing"
 	"time"
 
@@ -392,3 +394,107 @@ func TestRunBootstrap_DynamicLogitBias(t *testing.T) {
 		t.Errorf("user bias 999 bled into subagentClient: %v", subBiases)
 	}
 }
+
+// TestResolveSubagentTimeBudgetPrecedence guards the -subagent-timeout
+// resolution chain: an explicitly-passed flag beats the config.json
+// "subagent-timeout" entry, which beats the 24h default. Each case parses
+// real CLI args into a FlagSet registered exactly the way main() registers
+// the flag (default from appconfig.DefaultSubagentTimeout, single-sourced
+// usage string), so the explicit-set detection matches production.
+func TestResolveSubagentTimeBudgetPrecedence(t *testing.T) {
+	tests := []struct {
+		name                string
+		args                []string
+		cfg                 *appconfig.Config
+		want                time.Duration
+		wantWarningContains []string
+	}{
+		{
+			name: "nothing set defaults to 24h",
+			cfg:  nil,
+			want: 24 * time.Hour,
+		},
+		{
+			name: "config entry beats default",
+			args: nil,
+			cfg:  &appconfig.Config{SubagentTimeout: "90m"},
+			want: 90 * time.Minute,
+		},
+		{
+			name: "explicit flag beats config",
+			args: []string{"-subagent-timeout=2h"},
+			cfg:  &appconfig.Config{SubagentTimeout: "90m"},
+			want: 2 * time.Hour,
+		},
+		{
+			name: "explicit flag zero beats config and means unlimited",
+			args: []string{"-subagent-timeout=0"},
+			cfg:  &appconfig.Config{SubagentTimeout: "90m"},
+			want: 0,
+		},
+		{
+			name: "invalid config warns and falls back to default",
+			cfg:  &appconfig.Config{SubagentTimeout: "garbage"},
+			want: 24 * time.Hour,
+			wantWarningContains: []string{
+				"invalid",
+				"garbage",
+				"subagent-timeout",
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// Register -subagent-timeout the same way main() does.
+			fs := flag.NewFlagSet("subagent-timeout-precedence", flag.ContinueOnError)
+			flagValue := fs.Duration("subagent-timeout", appconfig.DefaultSubagentTimeout, subagentTimeoutUsage)
+			if err := fs.Parse(tt.args); err != nil {
+				t.Fatalf("Parse(%v): %v", tt.args, err)
+			}
+
+			got, warning := resolveSubagentTimeBudget(fs, tt.cfg, *flagValue)
+
+			if got != tt.want {
+				t.Fatalf("resolveSubagentTimeBudget() = %v, want %v", got, tt.want)
+			}
+			if len(tt.wantWarningContains) == 0 {
+				if warning != "" {
+					t.Fatalf("resolveSubagentTimeBudget() warning = %q, want empty", warning)
+				}
+				return
+			}
+			if warning == "" {
+				t.Fatal("resolveSubagentTimeBudget() warning is empty, want a warning")
+			}
+			for _, substring := range tt.wantWarningContains {
+				if !strings.Contains(warning, substring) {
+					t.Fatalf("resolveSubagentTimeBudget() warning = %q, want it to contain %q", warning, substring)
+				}
+			}
+		})
+	}
+}
+
+// TestSubagentTimeoutFlagDefaultAndUsage guards the -subagent-timeout flag
+// surface: the default must stay the 24h appconfig.DefaultSubagentTimeout
+// (long autonomous runs were the motivating use case for the budget) and the
+// usage string must keep documenting "0 = unlimited".
+func TestSubagentTimeoutFlagDefaultAndUsage(t *testing.T) {
+	if appconfig.DefaultSubagentTimeout != 24*time.Hour {
+		t.Fatalf("appconfig.DefaultSubagentTimeout = %v, want 24h", appconfig.DefaultSubagentTimeout)
+	}
+	if !strings.Contains(subagentTimeoutUsage, "0 = unlimited") {
+		t.Errorf("subagentTimeoutUsage = %q, want it to mention \"0 = unlimited\"", subagentTimeoutUsage)
+	}
+
+	fs := flag.NewFlagSet("subagent-timeout-usage", flag.ContinueOnError)
+	fs.Duration("subagent-timeout", appconfig.DefaultSubagentTimeout, subagentTimeoutUsage)
+	f := fs.Lookup("subagent-timeout")
+	if f == nil {
+		t.Fatal("flag -subagent-timeout was not registered")
+	}
+	if f.DefValue != "24h0m0s" {
+		t.Errorf("flag -subagent-timeout default = %q, want \"24h0m0s\"", f.DefValue)
+	}
+}