diff --git a/cmd/late/main.go b/cmd/late/main.go index d312724a..38122f1f 100644 --- a/cmd/late/main.go +++ b/cmd/late/main.go @@ -85,6 +85,9 @@ func main() { 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, "Max wall-clock time for one subagent run (0 = unlimited)") + subagentIdleTimeout := flag.Duration("subagent-idle-timeout", 15*time.Minute, "Notify when a subagent has been truly idle (no stream progress, no in-flight tool, no nested spawn) for this long (0 = off)") + subagentIdleKillAfter := flag.Duration("subagent-idle-kill-after", 0, "Kill a subagent that stays truly idle past this duration (0 = notify only)") 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") @@ -125,6 +128,13 @@ func main() { } flag.Parse() + // Record which flags were explicitly passed on the command line. The app + // config loads AFTER flag.Parse below, so flag.Visit (which reports only + // command-line-set flags) is the only reliable "explicit flag > config" + // precedence signal for resolvers such as ResolveSubagentTimeout. + explicitFlags := map[string]bool{} + flag.Visit(func(f *flag.Flag) { explicitFlags[f.Name] = true }) + tool.SetSqzEnabled(*enableSqzReq) if *versionReq { @@ -340,6 +350,16 @@ func main() { } } + // Resolve the global subagent run budget with precedence + // explicit --subagent-timeout flag > config.json "subagent_timeout" > + // DefaultSubagentTimeout (24h). A non-positive resolved budget ("0" or + // negative) means unlimited; an invalid config value is ignored with a + // warning and the default applies instead. + resolvedSubagentTimeout, subagentTimeoutWarning := appconfig.ResolveSubagentTimeout(appConfig, explicitFlags["subagent-timeout"], *subagentTimeout) + if subagentTimeoutWarning != "" { + fmt.Fprintf(os.Stderr, "Warning: %s\n", subagentTimeoutWarning) + } + // Parse explicit user logit bias overrides if provided var explicitUserLogitBias map[string]int if *logitBiasReq != "" { @@ -524,6 +544,10 @@ func main() { // Create root orchestrator // We'll add middlewares later once the program is started rootAgent := orchestrator.NewBaseOrchestrator(common.MainAgentID, sess, nil, 0) + // Idle watchdog policy applies to the root agent too: an orchestrator + // stuck with no stream progress, tool, or nested spawn reports idle (and, + // with --subagent-idle-kill-after, cancels its own run). + rootAgent.SetIdlePolicy(*subagentIdleTimeout, *subagentIdleKillAfter) model := tui.NewModel(rootAgent, renderer, appConfig) model.SetActiveThemeStyles(themeBytes) @@ -674,7 +698,24 @@ func main() { }() if *enableSubagentsReq { - runner := func(ctx context.Context, goal string, ctxFiles []string, agentType string) (string, error) { + runner := func(ctx context.Context, goal string, ctxFiles []string, agentType string, timeoutOverride *time.Duration) (string, error) { + // Effective wall-clock budget for this run. Context layering: + // parent ctx (cancellation) ⊇ run budget (deadline) — runCtx is + // derived from ctx, so cancelling the parent still cancels the + // child while the budget only adds a deadline. Precedence: + // per-spawn override when positive > global resolved budget; an + // explicit per-spawn "0" (unlimited) suppresses the global + // budget; an absent override falls back to the global value. + // runBudget is kept in a local var so error classification can + // report e.g. "time budget exhausted (2h)". + runBudget := effectiveSubagentBudget(timeoutOverride, resolvedSubagentTimeout) + runCtx := ctx + var runCancel context.CancelFunc = func() {} + if runBudget > 0 { + runCtx, runCancel = context.WithTimeout(ctx, runBudget) + } + defer runCancel() + var currentSubagentClient *client.Client if appConfig != nil { if setting, ok := appConfig.GetModelForAgent(agentType); ok { @@ -703,6 +744,52 @@ func main() { } child.SetMiddlewares(buildMiddlewares(pluginManager, p, child.Registry())) + // NewSubagentOrchestrator already set the child's context from + // the parent (agent.go: child.SetContext(parent.Context())) — + // keep that inheritance and layer the run budget on top: runCtx + // is derived from the runner ctx, so the parent ctx (cancellation) + // subsumes the budget (deadline). BaseOrchestrator.Execute then + // derives its run context from this one, so the deadline reaches + // the executor run loop. + if sa, ok := child.(interface{ SetContext(context.Context) }); ok { + sa.SetContext(runCtx) + } + + // The child runs its own idle watchdog with the same policy as + // the parent (it is a BaseOrchestrator too), so a stuck nested + // agent reports idle — or kills itself — independently. + if sa, ok := child.(interface { + SetIdlePolicy(idle, killAfter time.Duration) + }); ok { + sa.SetIdlePolicy(*subagentIdleTimeout, *subagentIdleKillAfter) + } + + // The child streams on its own session, so the parent shows no + // progress while the nested run executes. Keep the parent's + // activity alive with a 1/minute heartbeat and mark the nested + // spawn busy, so the parent's idle watchdog neither fires nor + // idle-kills while its child is legitimately working. + done := make(chan struct{}) + defer close(done) + go func() { + t := time.NewTicker(time.Minute) + defer t.Stop() + for { + select { + case <-done: + return + case <-t.C: + // rootAgent is a *orchestrator.BaseOrchestrator, + // which satisfies common.ActivityMarker; the parent + // stays visibly active while the child works. + rootAgent.MarkActivity() + } + } + }() + // Mark the nested spawn busy on the parent for the idle watchdog. + rootAgent.BeginNestedSpawn() + defer rootAgent.EndNestedSpawn() + res, err := child.Execute("") if err != nil { return "", err @@ -726,6 +813,21 @@ func main() { } } +// effectiveSubagentBudget selects the wall-clock budget for one subagent run. +// Precedence: a positive per-spawn override wins; an explicit per-spawn "0" +// (unlimited) suppresses the global budget; an absent override falls back to +// the global budget. Any non-positive result means unlimited (no deadline) — +// callers guard with "> 0" before deriving a WithTimeout context. +func effectiveSubagentBudget(timeoutOverride *time.Duration, globalBudget time.Duration) time.Duration { + if timeoutOverride != nil { + if *timeoutOverride > 0 { + return *timeoutOverride + } + return 0 // explicit per-spawn unlimited suppresses the global budget + } + return globalBudget +} + // deriveEffectiveSessionID derives this run's session ID from the FINAL // history path so resumed sessions keep their original ID. It returns "" // for empty or unsafe results (a crafted meta file could claim an ID like diff --git a/cmd/late/main_test.go b/cmd/late/main_test.go index 3b89c29c..19564f84 100644 --- a/cmd/late/main_test.go +++ b/cmd/late/main_test.go @@ -2,6 +2,7 @@ package main import ( "encoding/json" + "flag" "net/http" "net/http/httptest" "os" @@ -17,6 +18,100 @@ import ( "late/internal/session" ) +// TestEffectiveSubagentBudget guards the runner's budget precedence: +// a positive per-spawn override wins; an explicit per-spawn "0" (unlimited) +// suppresses the global budget; an absent override falls back to the global +// budget; any non-positive result means unlimited. +func TestEffectiveSubagentBudget(t *testing.T) { + positive := 2 * time.Hour + unlimited := time.Duration(0) + global := 24 * time.Hour + + tests := []struct { + name string + timeoutOverride *time.Duration + globalBudget time.Duration + want time.Duration + }{ + { + name: "absent override falls back to the global budget", + timeoutOverride: nil, + globalBudget: global, + want: global, + }, + { + name: "positive per-spawn override wins over the global budget", + timeoutOverride: &positive, + globalBudget: global, + want: positive, + }, + { + name: "explicit per-spawn unlimited suppresses the global budget", + timeoutOverride: &unlimited, + globalBudget: global, + want: 0, + }, + { + name: "per-spawn negative is unlimited", + timeoutOverride: &unlimited, // tool normalizes negatives to 0 before this point + globalBudget: global, + want: 0, + }, + { + name: "absent override with unlimited global stays unlimited", + timeoutOverride: nil, + globalBudget: 0, + want: 0, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := effectiveSubagentBudget(tt.timeoutOverride, tt.globalBudget); got != tt.want { + t.Fatalf("effectiveSubagentBudget(%v, %v) = %v, want %v", tt.timeoutOverride, tt.globalBudget, got, tt.want) + } + }) + } +} + +// TestSubagentTimeoutFlagVisitDetection guards the precedence mechanism used +// in main(): config.json loads AFTER flag.Parse, so "explicit flag > config" +// resolution relies on flag.Visit reporting only command-line-set flags. A +// Duration flag must be reported when passed — even when its value equals the +// default — and never reported when absent. +func TestSubagentTimeoutFlagVisitDetection(t *testing.T) { + parseExplicit := func(t *testing.T, args []string) map[string]bool { + t.Helper() + fs := flag.NewFlagSet("late", flag.ContinueOnError) + fs.Duration("subagent-timeout", appconfig.DefaultSubagentTimeout, "Max wall-clock time for one subagent run (0 = unlimited)") + if err := fs.Parse(args); err != nil { + t.Fatalf("Parse(%v): %v", args, err) + } + explicit := map[string]bool{} + fs.Visit(func(f *flag.Flag) { explicit[f.Name] = true }) + return explicit + } + + tests := []struct { + name string + args []string + want bool + }{ + {name: "flag passed is reported", args: []string{"-subagent-timeout=2h"}, want: true}, + {name: "flag absent is not reported", args: nil, want: false}, + {name: "flag passed at the default value is still reported", args: []string{"-subagent-timeout=24h"}, want: true}, + {name: "flag passed as zero is still reported", args: []string{"-subagent-timeout=0"}, want: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := parseExplicit(t, tt.args)["subagent-timeout"]; got != tt.want { + t.Fatalf(`explicit["subagent-timeout"] = %v, want %v`, got, tt.want) + } + }) + } +} + // TestPluginInlineTool_RequiresConfirmation guards the documented contract // that plugin inline tools (arbitrary scripts) go through the normal user // confirmation flow. plugin-example.md: "user confirmation still prompts diff --git a/docs/quickstart.md b/docs/quickstart.md index 1259a580..3a50174a 100644 --- a/docs/quickstart.md +++ b/docs/quickstart.md @@ -157,6 +157,22 @@ Approvals decay over time rather than becoming permanent trust forever. --- +## Subagent Control + +Subagents run under wall-clock budgets and an idle watchdog: + +| Flag | Default | Behavior | +| --- | --- | --- | +| `--subagent-timeout ` | 24h | Max wall-clock time for one subagent run (`0` = unlimited). Also settable as `subagent_timeout` in `config.json`; the flag wins over the config value. | +| `--subagent-idle-timeout ` | 15m | Notify when a subagent has been truly idle — no stream progress, no in-flight tool, no nested spawn (`0` = off). | +| `--subagent-idle-kill-after ` | 0 | Kill a subagent that stays truly idle past this duration (`0` = notify only). | + +The orchestrator can also budget a single run: `spawn_subagent` accepts an optional `timeout` argument (e.g. `"45m"`, `"2h"`; `"0"` = unlimited; omitted = the global value). + +Every `bash` call accepts an optional per-call `timeout` argument (e.g. `"30m"`; `"0"` = unlimited; omitted = the global default of 10m). Timed-out processes are killed with their whole process group, so runaway pipes and grandchildren cannot hang the session. + +--- + ## Run Fully Autonomously with Podman For unattended work, large refactors, or overnight runs, use `late-podman`. diff --git a/internal/common/interfaces.go b/internal/common/interfaces.go index f9c1dfca..dbcf85e8 100644 --- a/internal/common/interfaces.go +++ b/internal/common/interfaces.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "late/internal/client" + "time" ) // ToolRunner defines the functional signature for executing a single tool call. @@ -91,6 +92,25 @@ type MessageQueuedEvent struct { func (e MessageQueuedEvent) OrchestratorID() string { return e.ID } +// SubagentIdleEvent is sent when an agent has been truly idle — no stream +// progress, no in-flight tool, no in-flight nested subagent — for longer +// than the configured idle threshold. Probe carries the last few transcript +// entries so the recipient can decide whether the agent is stuck. +type SubagentIdleEvent struct { + ID string + IdleFor time.Duration + Probe []string +} + +func (e SubagentIdleEvent) OrchestratorID() string { return e.ID } + +// ActivityMarker is implemented by orchestrators that track their own +// activity for the idle watchdog. Nested subagent runners call MarkActivity +// on the parent so it does not look idle while its child is working. +type ActivityMarker interface { + MarkActivity() +} + // PromptRequest defines a generic requirement for user input. type PromptRequest struct { ID string diff --git a/internal/config/config.go b/internal/config/config.go index 1d165193..64f1a717 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -7,10 +7,16 @@ import ( "os" "path/filepath" "runtime" + "time" ) const DefaultOpenAIBaseURL = "http://localhost:8080" +// DefaultSubagentTimeout is the default wall-clock budget for a single +// subagent run, applied when neither the --subagent-timeout flag nor the +// config.json "subagent_timeout" entry provides a value. +const DefaultSubagentTimeout = 24 * time.Hour + type EnvLookup func(string) (string, bool) type OpenAISettings struct { @@ -61,6 +67,13 @@ type Config struct { // Enable via config file or the --save-subagent-histories CLI flag. SaveSubagentHistories bool `json:"save_subagent_histories,omitempty"` + // SubagentTimeout is an optional wall-clock budget for a single subagent + // run, parsed with time.ParseDuration (e.g. "45m", "2h"). "0" or a + // negative value means unlimited. Empty means the built-in default + // (DefaultSubagentTimeout). Precedence: an explicitly passed + // --subagent-timeout flag wins over this entry. + SubagentTimeout string `json:"subagent_timeout,omitempty"` + // Legacy subagent fields for backward compatibility SubagentBaseURL string `json:"subagent_base_url,omitempty"` SubagentAPIKey string `json:"subagent_api_key,omitempty"` @@ -241,6 +254,28 @@ func ResolveSaveSubagentHistories(cfg *Config, cliExplicit bool, cliValue bool, return false } +// ResolveSubagentTimeout resolves the global wall-clock budget for a single +// subagent run. Precedence: explicitly passed CLI flag > config.json +// "subagent_timeout" entry > DefaultSubagentTimeout (24h). +// +// A config value that parses via time.ParseDuration passes through as-is: +// "0" or a negative value means unlimited — callers guard with "> 0" and +// treat any non-positive budget as unlimited. An unparseable non-empty +// config value is ignored: the default is returned together with a warning +// for the caller to surface. +func ResolveSubagentTimeout(cfg *Config, cliExplicit bool, cliValue time.Duration) (time.Duration, string) { + if cliExplicit { + return cliValue, "" + } + if cfg != nil && cfg.SubagentTimeout != "" { + if parsed, err := time.ParseDuration(cfg.SubagentTimeout); err == nil { + return parsed, "" + } + return DefaultSubagentTimeout, fmt.Sprintf("ignoring invalid config.json subagent_timeout %q; using default 24h", cfg.SubagentTimeout) + } + return DefaultSubagentTimeout, "" +} + func nonEmptyEnv(lookup EnvLookup, key string) (string, bool) { if lookup == nil { return "", false diff --git a/internal/config/config_test.go b/internal/config/config_test.go index aa810d16..492ab6bd 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -1,10 +1,13 @@ package config import ( + "encoding/json" "os" "path/filepath" "runtime" + "strings" "testing" + "time" ) func TestLoadConfig_MissingFileCreatesDefault(t *testing.T) { @@ -599,6 +602,116 @@ func TestResolveSaveSubagentHistories(t *testing.T) { } } +func TestResolveSubagentTimeout(t *testing.T) { + tests := []struct { + name string + cfg *Config + cliExplicit bool + cliValue time.Duration + want time.Duration + wantWarning string + }{ + { + name: "explicit flag wins over config", + cfg: &Config{SubagentTimeout: "1h"}, + cliExplicit: true, + cliValue: 2 * time.Minute, + want: 2 * time.Minute, + }, + { + name: "explicit zero flag wins over config (unlimited)", + cfg: &Config{SubagentTimeout: "1h"}, + cliExplicit: true, + cliValue: 0, + want: 0, + }, + { + name: "config parses to the configured budget", + cfg: &Config{SubagentTimeout: "24h"}, + want: 24 * time.Hour, + }, + { + name: "config minutes parse", + cfg: &Config{SubagentTimeout: "45m"}, + want: 45 * time.Minute, + }, + { + name: "config zero passes through as unlimited", + cfg: &Config{SubagentTimeout: "0"}, + want: 0, + }, + { + name: "config negative passes through as unlimited", + cfg: &Config{SubagentTimeout: "-5m"}, + want: -5 * time.Minute, + }, + { + name: "invalid config warns and falls back to default", + cfg: &Config{SubagentTimeout: "garbage"}, + want: DefaultSubagentTimeout, + wantWarning: `ignoring invalid config.json subagent_timeout "garbage"; using default 24h`, + }, + { + name: "unitless config value warns and falls back to default", + cfg: &Config{SubagentTimeout: "5"}, + want: DefaultSubagentTimeout, + wantWarning: `ignoring invalid config.json subagent_timeout "5"; using default 24h`, + }, + { + name: "empty config entry uses default", + cfg: &Config{}, + want: DefaultSubagentTimeout, + }, + { + name: "nil config uses default", + cfg: nil, + want: DefaultSubagentTimeout, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, warning := ResolveSubagentTimeout(tt.cfg, tt.cliExplicit, tt.cliValue) + if got != tt.want { + t.Fatalf("ResolveSubagentTimeout() = %v, want %v", got, tt.want) + } + if warning != tt.wantWarning { + t.Fatalf("ResolveSubagentTimeout() warning = %q, want %q", warning, tt.wantWarning) + } + }) + } +} + +// TestConfig_SubagentTimeoutJSONRoundTrip pins the config.json key spelling +// and round-trip behavior of the subagent_timeout entry: omitting it stays +// omitempty, and a set value survives Marshal/Unmarshal unchanged. The key +// is snake_case, matching the sibling entries in the config schema. +func TestConfig_SubagentTimeoutJSONRoundTrip(t *testing.T) { + data, err := json.Marshal(&Config{}) + if err != nil { + t.Fatalf("Marshal() error = %v", err) + } + if strings.Contains(string(data), "subagent_timeout") { + t.Fatalf("empty SubagentTimeout must be omitted, got %s", data) + } + + data, err = json.Marshal(&Config{SubagentTimeout: "45m"}) + if err != nil { + t.Fatalf("Marshal() error = %v", err) + } + if !strings.Contains(string(data), `"subagent_timeout":"45m"`) { + t.Fatalf("expected subagent_timeout key in JSON, got %s", data) + } + + var back Config + if err := json.Unmarshal(data, &back); err != nil { + t.Fatalf("Unmarshal() error = %v", err) + } + if back.SubagentTimeout != "45m" { + t.Fatalf("round-tripped SubagentTimeout = %q, want %q", back.SubagentTimeout, "45m") + } +} + func TestConfig_GetModelForAgent(t *testing.T) { cfg := &Config{ Models: []ModelSetting{ diff --git a/internal/executor/executor.go b/internal/executor/executor.go index eb9be73d..f4412a9d 100644 --- a/internal/executor/executor.go +++ b/internal/executor/executor.go @@ -101,9 +101,36 @@ func ExecuteToolCalls(ctx context.Context, sess *session.Session, toolCalls []cl } } - result, err := runner(ctx, tc) + // Each tool call runs in its own cancellable context so a hung tool + // can be killed individually (via sess.CancelInFlightTool — used by + // the orchestrator's idle watchdog, or the user) without aborting the + // whole run. The cancel func is registered on the session for the + // duration of the call; every iteration derives a FRESH toolCtx from + // the parent ctx, so a cancelled call never poisons its successors. + toolCtx, toolCancel := context.WithCancel(ctx) + sess.SetInFlightToolCancel(toolCancel) + result, err := runner(toolCtx, tc) + sess.ClearInFlightToolCancel() + toolCancel() // the call is done; release the derived context + if err != nil { - result = fmt.Sprintf("Error executing tool %s: %v", tc.Function.Name, err) + if ctx.Err() == nil && toolCtx.Err() != nil { + // The tool call's own context was cancelled while the parent + // run is still alive: this was an in-flight-tool kill (harness + // idle watchdog or user), not a stop request. Surface the + // cancellation as a normal tool result — the model sees it and + // can recover — and keep processing the remaining calls. + // (A per-call timeout would land here too, as + // toolCtx.Err() == context.DeadlineExceeded; none exists on + // this branch yet — the shell timeout is enforced inside the + // tool itself and returns a normal result.) + result = "tool cancelled by the harness idle watchdog" + } else { + // Parent cancelled (or a plain tool failure): keep today's + // behaviour of noting the error; the run loop's subsequent + // ctx checks unwind the run. + result = fmt.Sprintf("Error executing tool %s: %v", tc.Function.Name, err) + } } if err := sess.AddToolResultMessage(tc.ID, result); err != nil { return err diff --git a/internal/executor/inflight_test.go b/internal/executor/inflight_test.go new file mode 100644 index 00000000..596bc814 --- /dev/null +++ b/internal/executor/inflight_test.go @@ -0,0 +1,192 @@ +package executor + +import ( + "context" + "encoding/json" + "fmt" + "strings" + "sync/atomic" + "testing" + "time" + + "late/internal/client" + "late/internal/session" +) + +// fakeTool is a minimal common.Tool implementation for executor tests. +type fakeTool struct { + name string + exec func(ctx context.Context, args json.RawMessage) (string, error) +} + +func (f *fakeTool) Name() string { return f.name } +func (f *fakeTool) Description() string { return "fake tool for executor tests" } +func (f *fakeTool) Parameters() json.RawMessage { + return json.RawMessage(`{"type":"object"}`) +} +func (f *fakeTool) Execute(ctx context.Context, args json.RawMessage) (string, error) { + return f.exec(ctx, args) +} +func (f *fakeTool) RequiresConfirmation(json.RawMessage) bool { return false } +func (f *fakeTool) CallString(json.RawMessage) string { return f.name } + +// TestExecuteToolCalls_InFlightToolKillContinues verifies the killable +// in-flight tool path: cancelling the in-flight tool (via the session hook the +// orchestrator's idle watchdog uses) turns the hung call into a "tool +// cancelled" result, ExecuteToolCalls keeps going, and the NEXT tool call runs +// with a fresh context on a still-alive parent. +func TestExecuteToolCalls_InFlightToolKillContinues(t *testing.T) { + c := client.NewClient(client.Config{BaseURL: "http://localhost:0"}) + sess := session.New(c, "", nil, "", true) // no history path: in-memory session + + var sawCancel atomic.Bool + started := make(chan struct{}) + hung := &fakeTool{name: "hung_tool", exec: func(ctx context.Context, args json.RawMessage) (string, error) { + close(started) + <-ctx.Done() + sawCancel.Store(true) + return "", fmt.Errorf("tool cancelled: %v", ctx.Err()) + }} + quick := &fakeTool{name: "quick_tool", exec: func(ctx context.Context, args json.RawMessage) (string, error) { + if err := ctx.Err(); err != nil { + return "", fmt.Errorf("inherited a cancelled context: %v", err) + } + return "second ok", nil + }} + sess.Registry.Register(hung) + sess.Registry.Register(quick) + + toolCalls := []client.ToolCall{ + {ID: "tc_1", Function: client.FunctionCall{Name: "hung_tool", Arguments: "{}"}}, + {ID: "tc_2", Function: client.FunctionCall{Name: "quick_tool", Arguments: "{}"}}, + } + + parentCtx, parentCancel := context.WithCancel(context.Background()) + defer parentCancel() + + done := make(chan error, 1) + go func() { + done <- ExecuteToolCalls(parentCtx, sess, toolCalls, nil) + }() + + // Wait until the hung tool is executing — the executor registered its + // cancel on the session before invoking it. + <-started + + // Kill the in-flight tool exactly the way the idle watchdog does. + if !sess.CancelInFlightTool() { + t.Fatal("expected a registered in-flight tool cancel while the hung tool ran") + } + + select { + case err := <-done: + if err != nil { + t.Fatalf("ExecuteToolCalls returned an error after a tool kill: %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("ExecuteToolCalls did not finish after the in-flight tool was killed") + } + + if !sawCancel.Load() { + t.Error("the hung tool did not observe the cancellation of its context") + } + if parentCtx.Err() != nil { + t.Errorf("parent context must stay alive after a tool kill, got %v", parentCtx.Err()) + } + + // Both results must be recorded: the cancellation for the killed call and + // the normal output for the call that ran afterwards. + if len(sess.History) != 2 { + t.Fatalf("expected 2 history entries, got %d", len(sess.History)) + } + first := sess.History[0].Content.String() + if !strings.Contains(first, "cancelled") { + t.Errorf("first result must report the tool cancellation, got %q", first) + } + if strings.Contains(first, "Error executing tool") { + t.Errorf("a tool kill is not a tool error, got %q", first) + } + if second := sess.History[1].Content.String(); second != "second ok" { + t.Errorf("second tool call must still have run with a fresh context, got %q", second) + } + + // The hook must be cleared once the calls finished. + if sess.CancelInFlightTool() { + t.Error("in-flight cancel hook was not cleared after ExecuteToolCalls returned") + } +} + +// TestExecuteToolCalls_ParentCancelStillAborts pins the pre-existing +// semantics for a cancelled PARENT context: the result notes the error (it is +// not rewritten into a tool-cancellation message) and the loop still processes +// the remaining calls — the run itself is unwound by the run loop's ctx +// checks, not by ExecuteToolCalls. +func TestExecuteToolCalls_ParentCancelStillAborts(t *testing.T) { + c := client.NewClient(client.Config{BaseURL: "http://localhost:0"}) + sess := session.New(c, "", nil, "", true) + + started := make(chan struct{}) + blocker := &fakeTool{name: "blocking_tool", exec: func(ctx context.Context, args json.RawMessage) (string, error) { + close(started) + <-ctx.Done() + return "", ctx.Err() + }} + probe := &fakeTool{name: "probe_tool", exec: func(ctx context.Context, args json.RawMessage) (string, error) { + if err := ctx.Err(); err != nil { + return "", err + } + return "probe ok", nil + }} + sess.Registry.Register(blocker) + sess.Registry.Register(probe) + + toolCalls := []client.ToolCall{ + {ID: "tc_1", Function: client.FunctionCall{Name: "blocking_tool", Arguments: "{}"}}, + {ID: "tc_2", Function: client.FunctionCall{Name: "probe_tool", Arguments: "{}"}}, + } + + parentCtx, parentCancel := context.WithCancel(context.Background()) + + done := make(chan error, 1) + go func() { + done <- ExecuteToolCalls(parentCtx, sess, toolCalls, nil) + }() + + <-started + parentCancel() // cancel the parent mid-first-call + + select { + case err := <-done: + if err != nil { + t.Fatalf("ExecuteToolCalls returned an error: %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("ExecuteToolCalls did not finish after the parent context was cancelled") + } + + if parentCtx.Err() != context.Canceled { + t.Fatalf("parent ctx.Err() = %v, want context.Canceled", parentCtx.Err()) + } + + if len(sess.History) != 2 { + t.Fatalf("expected 2 history entries, got %d", len(sess.History)) + } + first := sess.History[0].Content.String() + if !strings.Contains(first, "Error executing tool blocking_tool") || + !strings.Contains(first, "context canceled") { + t.Errorf("a cancelled parent must note the error, got %q", first) + } + if strings.Contains(first, "tool cancelled by the harness idle watchdog") { + t.Errorf("a parent cancellation must not be reported as a tool kill, got %q", first) + } + // Pre-change semantics: the loop kept processing the remaining calls; the + // probe tool fails fast on the dead context. + second := sess.History[1].Content.String() + if !strings.Contains(second, "Error executing tool probe_tool") { + t.Errorf("expected the second call to note the dead context, got %q", second) + } + + if sess.CancelInFlightTool() { + t.Error("in-flight cancel hook was not cleared after ExecuteToolCalls returned") + } +} diff --git a/internal/orchestrator/base.go b/internal/orchestrator/base.go index 7d6b8052..d7fafe50 100644 --- a/internal/orchestrator/base.go +++ b/internal/orchestrator/base.go @@ -13,6 +13,8 @@ import ( "path/filepath" "strings" "sync" + "sync/atomic" + "time" ) // BaseOrchestrator implements common.Orchestrator and manages an agent's run loop. @@ -41,14 +43,61 @@ type BaseOrchestrator struct { // Max turns configuration maxTurns int + + // Activity tracking for the idle watchdog. lastActivity holds unix-nano + // of the last observed sign of life (stream chunk, tool execution, + // nested-spawn heartbeat); inFlightTools and nestedSpawns count + // outstanding tool executions and nested subagent runs. All three are + // atomic; idleNotified enforces the once-per-idle-episode notification. + // oldestToolStartAt holds the unix-nano start time of the oldest + // in-flight tool (0 when none) so the busy check can tell a progressing + // tool (younger than the idle threshold) from one that has stalled past + // it — a stalled tool must not mask idleness forever, or the watchdog + // could never kill the hung tool and rescue the agent. + lastActivity atomic.Int64 + inFlightTools atomic.Int64 + oldestToolStartAt atomic.Int64 + nestedSpawns atomic.Int64 + idleNotified atomic.Bool + + // toolKillDone records that the idle watchdog already used its first + // escalation stage for this run (killed the hung in-flight tool), so the + // next sustained-idle tick goes straight to the agent-level kill. + // Once-per-stage semantics: stage 1 fires at most once per run, stage 2 + // (the agent kill) at most once because cancel() readies ctx.Done and + // stops the watchdog. + toolKillDone atomic.Bool + + // Idle-watchdog policy, guarded by mu and set via SetIdlePolicy before a + // run starts. idleTimeout <= 0 disables the watchdog; idleKillAfter <= 0 + // means notify only. idleTickInterval defaults to defaultIdleTickInterval + // and is overridable (same package) so tests can drive the watchdog + // deterministically. + idleTimeout time.Duration + idleKillAfter time.Duration + idleTickInterval time.Duration + + // idleKillReason records why the idle watchdog cancelled this run; guarded + // by mu. Empty unless the watchdog killed the run. + idleKillReason string } +// Idle-watchdog tuning. defaultIdleTickInterval is how often the watchdog +// re-checks for idleness; idleProbeLines is how many transcript entries are +// rendered into a SubagentIdleEvent probe; idleProbeLineMaxLen caps each +// rendered probe line. +const ( + defaultIdleTickInterval = 30 * time.Second + idleProbeLines = 3 + idleProbeLineMaxLen = 160 +) + func NewBaseOrchestrator(id string, sess *session.Session, middlewares []common.ToolMiddleware, maxTurns int) *BaseOrchestrator { childSeq := 0 if sess != nil { childSeq = sess.SubagentSeq() } - return &BaseOrchestrator{ + o := &BaseOrchestrator{ id: id, sess: sess, middlewares: middlewares, @@ -58,6 +107,10 @@ func NewBaseOrchestrator(id string, sess *session.Session, middlewares []common. maxTurns: maxTurns, childSeq: childSeq, } + // Construction counts as activity so an immediate run starts with a + // fresh idle episode. + o.lastActivity.Store(time.Now().UnixNano()) + return o } func (o *BaseOrchestrator) SetMiddlewares(middlewares []common.ToolMiddleware) { @@ -87,6 +140,273 @@ func (o *BaseOrchestrator) SetMaxTurns(maxTurns int) { o.maxTurns = maxTurns } +// SetIdlePolicy configures the idle watchdog for this orchestrator: idle is +// the "truly idle" threshold (0 = watchdog off) and killAfter is the +// sustained-idle point at which the orchestrator cancels its own run +// (0 = notify only). Must be called before Execute/Submit to apply to a run. +func (o *BaseOrchestrator) SetIdlePolicy(idle, killAfter time.Duration) { + o.mu.Lock() + defer o.mu.Unlock() + o.idleTimeout = idle + o.idleKillAfter = killAfter +} + +// MarkActivity records a sign of life — a streamed chunk, a tool execution, +// or a nested-spawn heartbeat — and re-arms the idle notification so a new +// idle episode can be reported after activity resumes. It implements +// common.ActivityMarker. +func (o *BaseOrchestrator) MarkActivity() { + o.lastActivity.Store(time.Now().UnixNano()) + o.idleNotified.Store(false) +} + +// BeginNestedSpawn marks a nested subagent run as in flight so the idle +// watchdog treats this orchestrator as busy for the child's duration. +func (o *BaseOrchestrator) BeginNestedSpawn() { o.nestedSpawns.Add(1) } + +// EndNestedSpawn marks a previously begun nested subagent run as finished. +func (o *BaseOrchestrator) EndNestedSpawn() { o.nestedSpawns.Add(-1) } + +// IdleKillReason returns the recorded kill reason when the idle watchdog +// self-cancelled this run, or "" otherwise. +func (o *BaseOrchestrator) IdleKillReason() string { + o.mu.RLock() + defer o.mu.RUnlock() + return o.idleKillReason +} + +// activityMiddleware is the internal, always-present outermost middleware: +// every tool execution bumps activity and the in-flight tool counter, and +// stamps the in-flight start time used by the idle watchdog's busy check. It +// is inserted OUTERMOST (before user middlewares), so the confirmation +// middleware runs INSIDE this wrapper — a tool waiting for user approval +// still counts as in-flight, i.e. awaiting-approval counts as active. That +// is by design: a paused-for-approval agent is not idle. +func (o *BaseOrchestrator) activityMiddleware() common.ToolMiddleware { + return func(next common.ToolRunner) common.ToolRunner { + return func(ctx context.Context, tc client.ToolCall) (string, error) { + o.MarkActivity() + o.beginToolInFlight() + defer o.endToolInFlight() + return next(ctx, tc) + } + } +} + +// beginToolInFlight records the start of a tool execution. With concurrent +// tools it keeps the OLDEST start time — the busy check only needs the +// longest-running call — and the stamp is cleared again once no tool is +// in flight. +func (o *BaseOrchestrator) beginToolInFlight() { + o.inFlightTools.Add(1) + now := time.Now().UnixNano() + for { + prev := o.oldestToolStartAt.Load() + if prev != 0 && prev <= now { + return + } + if o.oldestToolStartAt.CompareAndSwap(prev, now) { + return + } + } +} + +// endToolInFlight marks a tool execution as finished and clears the oldest +// in-flight timestamp when the in-flight count drops back to zero. +func (o *BaseOrchestrator) endToolInFlight() { + if o.inFlightTools.Add(-1) == 0 { + o.oldestToolStartAt.Store(0) + } +} + +// withActivityMiddleware returns the middleware chain passed to RunLoop with +// the internal activity middleware prepended as the outermost layer. The +// chain is rebuilt fresh so the stored o.middlewares slice is never mutated. +func (o *BaseOrchestrator) withActivityMiddleware() []common.ToolMiddleware { + chain := make([]common.ToolMiddleware, 0, len(o.middlewares)+1) + chain = append(chain, o.activityMiddleware()) + chain = append(chain, o.middlewares...) + return chain +} + +// startIdleWatchdog launches the idle watchdog for one run; it stops when the +// run context is cancelled. On every tick it checks whether the orchestrator +// has been truly idle — no stream progress, no in-flight nested spawn, and no +// tool call younger than the idle threshold — for longer than the configured +// idle threshold. An in-flight tool counts as progress only while it is +// YOUNGER than the idle threshold: a tool stuck for longer than that stops +// masking idleness, which is what makes the tool-first kill (stage 1 of the +// escalation below) reachable for hung tools. The first qualifying tick emits +// one SubagentIdleEvent (once per idle episode; MarkActivity re-arms). When a +// kill threshold is configured, the kill escalates in two stages (see the +// kill branch below for why that check is not gated on the once-per-episode +// flag). +func (o *BaseOrchestrator) startIdleWatchdog(ctx context.Context, cancel context.CancelFunc) { + o.mu.RLock() + idleTimeout := o.idleTimeout + idleKillAfter := o.idleKillAfter + tickInterval := o.idleTickInterval + o.mu.RUnlock() + + if idleTimeout <= 0 { + return + } + if tickInterval <= 0 { + tickInterval = defaultIdleTickInterval + } + + go func() { + ticker := time.NewTicker(tickInterval) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + } + + idle := time.Since(time.Unix(0, o.lastActivity.Load())) + // Busy means real progress is still possible: a nested subagent + // run, or a tool call in flight for LESS than the idle threshold. + // A tool that outlives idleTimeout is considered stalled — it no + // longer suppresses idleness, so a hung tool cannot hide from the + // watchdog forever. + busy := o.nestedSpawns.Load() > 0 + if start := o.oldestToolStartAt.Load(); start != 0 && + time.Since(time.Unix(0, start)) < idleTimeout { + busy = true + } + if idle < idleTimeout || busy { + continue + } + + // Once per idle episode: Swap(false→true) fires only on the + // first qualifying tick after the last MarkActivity. + if !o.idleNotified.Swap(true) { + select { + case o.eventCh <- common.SubagentIdleEvent{ID: o.id, IdleFor: idle, Probe: o.idleProbe()}: + default: + } + } + + // Two-stage sustained-idle kill. Like the notification above, the + // threshold check runs on EVERY qualifying tick: the kill + // threshold normally sits well past the notify threshold (e.g. + // notify at 15m, kill at 30m), and gating it on the + // once-per-episode flag would swallow it — the notify tick fires + // long before idle reaches the kill threshold. + // + // Stage 1 (first qualifying tick that finds a stalled tool in + // flight): kill the TOOL, not the agent. executor.ExecuteToolCalls + // registers each call's cancellable context on the session, so + // CancelInFlightTool makes the hung call return a "tool cancelled" + // result and the run continues — the model sees the cancellation + // and can recover. The watchdog then re-checks on the NEXT tick: + // if activity resumed (next tool call, streamed chunk), nothing + // else happens; if idleness persists past the kill threshold, + // stage 2 fires. Once-per-stage: toolKillDone gates stage 1 to a + // single shot per run. + // + // Stage 2 (next qualifying tick, or the same one when no tool was + // in flight to kill): cancel our own run context — the + // orchestrator probed the transcript and decided this agent is + // stuck. cancel() readies ctx.Done, so stage 2 records and + // cancels at most once. + if idleKillAfter > 0 && idle >= idleKillAfter { + if start := o.oldestToolStartAt.Load(); start != 0 && !o.toolKillDone.Load() { + if o.sess != nil && o.sess.CancelInFlightTool() { + // Stage 1 done. Skip the agent kill this tick and + // re-evaluate on the next one. + o.toolKillDone.Store(true) + continue + } + // The tool finished between the busy check and the kill + // attempt — nothing left to cancel; fall through to + // stage 2. + } + + // Record the probe lines in the kill reason for the crash + // classification, noting an earlier tool kill that failed to + // rescue the run. + probe := o.idleProbe() + reason := fmt.Sprintf( + "idle for %s (kill threshold %s); last transcript entries: %s", + idle.Truncate(time.Second), idleKillAfter, strings.Join(probe, " | "), + ) + if o.toolKillDone.Load() { + reason = "in-flight tool cancelled; " + reason + } + o.mu.Lock() + o.idleKillReason = reason + o.mu.Unlock() + if cancel != nil { + cancel() + } + } + } + }() +} + +// idleProbe renders the current transcript tail for idle events and kill +// reasons. Best-effort: history is read without locking, matching +// BaseOrchestrator.History(); the run loop may append concurrently. +func (o *BaseOrchestrator) idleProbe() []string { + if o.sess == nil { + return nil + } + return lastTranscriptLines(o.sess.History, idleProbeLines) +} + +// lastTranscriptLines renders the last n history entries as short single-line +// strings ("role: first line") — the transcript probe carried by idle events +// so the recipient can judge whether the agent is stuck. Assistant messages +// that only contain tool calls fall back to the tool names. +func lastTranscriptLines(msgs []client.ChatMessage, n int) []string { + if n <= 0 || len(msgs) == 0 { + return nil + } + start := len(msgs) - n + if start < 0 { + start = 0 + } + lines := make([]string, 0, len(msgs)-start) + for _, msg := range msgs[start:] { + summary := firstLine(strings.TrimSpace(msg.Content.String())) + if summary == "" && len(msg.ToolCalls) > 0 { + names := make([]string, 0, len(msg.ToolCalls)) + for _, tc := range msg.ToolCalls { + names = append(names, tc.Function.Name) + } + summary = "called " + strings.Join(names, ", ") + } + if summary == "" { + continue + } + lines = append(lines, truncateRunes(fmt.Sprintf("%s: %s", msg.Role, summary), idleProbeLineMaxLen)) + } + return lines +} + +// firstLine returns s up to the first newline. +func firstLine(s string) string { + if idx := strings.IndexAny(s, "\r\n"); idx >= 0 { + return s[:idx] + } + return s +} + +// truncateRunes shortens s to max runes, appending "..." when truncated. +func truncateRunes(s string, max int) string { + runes := []rune(s) + if len(runes) <= max { + return s + } + if max <= 3 { + return string(runes[:max]) + } + return string(runes[:max-3]) + "..." +} + func (o *BaseOrchestrator) MaxTokens() int { o.mu.RLock() defer o.mu.RUnlock() @@ -205,6 +525,12 @@ func (o *BaseOrchestrator) Execute(text string) (string, error) { defer cancel() + // Fresh idle episode for this run: construction or the previous run's + // last activity must not leak into the watchdog's baseline. The watchdog + // stops when the deferred cancel above fires. + o.MarkActivity() + o.startIdleWatchdog(ctx, cancel) + // Inject orchestrator ID into context for tool interactions ctx = context.WithValue(ctx, common.OrchestratorIDKey, o.id) @@ -258,6 +584,8 @@ func (o *BaseOrchestrator) Execute(text string) (string, error) { onStartTurn, onEndTurn, func(res common.StreamResult) { + // Stream progress counts as activity for the idle watchdog. + o.MarkActivity() o.mu.Lock() o.acc.Append(res) accCopy := o.acc @@ -271,7 +599,7 @@ func (o *BaseOrchestrator) Execute(text string) (string, error) { Usage: accCopy.Usage, } }, - o.middlewares, + o.withActivityMiddleware(), ) if err != nil { @@ -292,6 +620,11 @@ func (o *BaseOrchestrator) run() { defer cancel() // Ensure we don't leak the context when run() finishes + // Fresh idle episode for this run; the watchdog stops when the deferred + // cancel above fires. + o.MarkActivity() + o.startIdleWatchdog(ctx, cancel) + // Inject orchestrator ID into context for tool interactions ctx = context.WithValue(ctx, common.OrchestratorIDKey, o.id) @@ -331,6 +664,8 @@ func (o *BaseOrchestrator) run() { onStartTurn, onEndTurn, func(res common.StreamResult) { + // Stream progress counts as activity for the idle watchdog. + o.MarkActivity() o.mu.Lock() o.acc.Append(res) accCopy := o.acc // Copy for event @@ -344,7 +679,7 @@ func (o *BaseOrchestrator) run() { Usage: accCopy.Usage, } }, - o.middlewares, + o.withActivityMiddleware(), ) // Reset accumulator after finished or ready for next turn diff --git a/internal/orchestrator/base_idle_test.go b/internal/orchestrator/base_idle_test.go new file mode 100644 index 00000000..c379796f --- /dev/null +++ b/internal/orchestrator/base_idle_test.go @@ -0,0 +1,453 @@ +package orchestrator + +import ( + "context" + "encoding/json" + "fmt" + "strings" + "testing" + "time" + + "late/internal/client" + "late/internal/common" + "late/internal/executor" + "late/internal/session" +) + +// The idle watchdog is driven deterministically by injecting a tiny idle +// threshold and tick interval (both unexported fields, same package). The +// production defaults — a 30s tick and a threshold supplied by the +// --subagent-idle-timeout flag (15m, 0 = off) — are untouched. + +// newIdleTestOrchestrator returns an orchestrator with the given idle policy +// and tick interval, seeded with the provided history. lastActivity is +// stamped at construction, matching NewBaseOrchestrator. +func newIdleTestOrchestrator(t *testing.T, idle, killAfter, tick time.Duration, history []client.ChatMessage) *BaseOrchestrator { + t.Helper() + sess := session.New(nil, "", history, "", false) + o := NewBaseOrchestrator("idle-test", sess, nil, 10) + o.SetIdlePolicy(idle, killAfter) + o.idleTickInterval = tick + return o +} + +// collectIdleEvents drains the orchestrator's event channel for the given +// duration and returns every SubagentIdleEvent received in that window. +func collectIdleEvents(t *testing.T, o *BaseOrchestrator, wait time.Duration) []common.SubagentIdleEvent { + t.Helper() + var events []common.SubagentIdleEvent + deadline := time.After(wait) + for { + select { + case ev := <-o.eventCh: + if idle, ok := ev.(common.SubagentIdleEvent); ok { + events = append(events, idle) + } + case <-deadline: + return events + } + } +} + +// waitFor polls cond until it holds or the timeout elapses. +func waitFor(t *testing.T, timeout time.Duration, cond func() bool) { + t.Helper() + deadline := time.Now().Add(timeout) + for time.Now().Before(deadline) { + if cond() { + return + } + time.Sleep(2 * time.Millisecond) + } + t.Fatal("condition not reached within timeout") +} + +// TestIdleWatchdogFiresAfterSilence verifies the watchdog emits exactly one +// SubagentIdleEvent per idle episode: one event once the silence passes the +// threshold, no duplicates while the silence persists, no event right after +// activity resumes, and a fresh event once the agent goes idle again. The +// probe must carry the last transcript entries only. +func TestIdleWatchdogFiresAfterSilence(t *testing.T) { + history := []client.ChatMessage{ + {Role: "user", Content: client.TextContent("first task")}, + {Role: "assistant", Content: client.TextContent("working on it")}, + {Role: "assistant", Content: client.TextContent("thinking deeply\nabout options")}, + {Role: "user", Content: client.TextContent("continue")}, + } + o := newIdleTestOrchestrator(t, 60*time.Millisecond, 0, 10*time.Millisecond, history) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + o.startIdleWatchdog(ctx, cancel) + + // First idle episode: exactly one event once the threshold passes. + events := collectIdleEvents(t, o, 250*time.Millisecond) + if len(events) != 1 { + t.Fatalf("expected exactly 1 idle event per episode, got %d", len(events)) + } + if events[0].ID != o.id { + t.Errorf("event ID = %q, want %q", events[0].ID, o.id) + } + if events[0].IdleFor <= 0 { + t.Errorf("IdleFor = %v, want > 0", events[0].IdleFor) + } + + // Probe: the last 3 history entries, rendered as "role: first line". + wantProbe := []string{ + "assistant: working on it", + "assistant: thinking deeply", // first line only + "user: continue", + } + if len(events[0].Probe) != len(wantProbe) { + t.Fatalf("probe length = %d (%v), want %d", len(events[0].Probe), events[0].Probe, len(wantProbe)) + } + for i, want := range wantProbe { + if events[0].Probe[i] != want { + t.Errorf("probe[%d] = %q, want %q", i, events[0].Probe[i], want) + } + } + if strings.Contains(strings.Join(events[0].Probe, "\n"), "first task") { + t.Errorf("probe must only contain the last entries, got %v", events[0].Probe) + } + + // Idle persists: once-per-episode means no further events. + if again := collectIdleEvents(t, o, 60*time.Millisecond); len(again) != 0 { + t.Fatalf("idle event re-emitted while the episode persisted: %v", again) + } + + // Re-arm: activity resets the episode; no event while still fresh. + o.MarkActivity() + if again := collectIdleEvents(t, o, 30*time.Millisecond); len(again) != 0 { + t.Fatalf("idle event fired right after activity resumed: %v", again) + } + + // Idle again: a new episode produces exactly one more event. + if again := collectIdleEvents(t, o, 250*time.Millisecond); len(again) != 1 { + t.Fatalf("expected exactly 1 re-armed idle event, got %d", len(again)) + } +} + +// TestIdleWatchdogSuppressedWhileToolInFlight verifies that an in-flight tool +// that completes just before the idle threshold still counts as progress: the +// idle event is suppressed while it runs, and fires once the accumulated +// silence (measured from the tool's start) passes the threshold. A tool that +// outlives the threshold behaves differently — see +// TestIdleWatchdogKillsHungToolBeforeAgent. +func TestIdleWatchdogSuppressedWhileToolInFlight(t *testing.T) { + o := newIdleTestOrchestrator(t, 100*time.Millisecond, 0, 10*time.Millisecond, []client.ChatMessage{ + {Role: "user", Content: client.TextContent("goal")}, + }) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + o.startIdleWatchdog(ctx, cancel) + + started := make(chan struct{}) + finished := make(chan struct{}) + base := func(ctx context.Context, tc client.ToolCall) (string, error) { + close(started) + time.Sleep(60 * time.Millisecond) // completes before the 100ms threshold + return "done", nil + } + toolRunner := o.activityMiddleware()(base) + + go func() { + defer close(finished) + _, _ = toolRunner(ctx, client.ToolCall{ + Type: "function", + Function: client.FunctionCall{Name: "fake_slow_tool", Arguments: "{}"}, + }) + }() + + <-started + waitFor(t, 2*time.Second, func() bool { return o.inFlightTools.Load() == 1 }) + + // The young in-flight tool suppresses the idle event. + if events := collectIdleEvents(t, o, 40*time.Millisecond); len(events) != 0 { + t.Fatalf("idle event fired while a young tool was in flight: %v", events) + } + + <-finished + waitFor(t, 2*time.Second, func() bool { return o.inFlightTools.Load() == 0 }) + waitFor(t, 2*time.Second, func() bool { return o.oldestToolStartAt.Load() == 0 }) + + // Watchdog is alive: with the tool finished and no new activity, the + // episode fires exactly once. + if events := collectIdleEvents(t, o, 250*time.Millisecond); len(events) != 1 { + t.Fatalf("expected exactly 1 idle event after the tool finished, got %d", len(events)) + } +} + +// fakeIdleTool is a minimal common.Tool implementation for idle-watchdog +// tests, registered on the test session so executor.ExecuteToolCalls can run +// it through the orchestrator's activity middleware. +type fakeIdleTool struct { + name string + exec func(ctx context.Context, args json.RawMessage) (string, error) +} + +func (f *fakeIdleTool) Name() string { return f.name } +func (f *fakeIdleTool) Description() string { return "fake tool for idle tests" } +func (f *fakeIdleTool) Parameters() json.RawMessage { + return json.RawMessage(`{"type":"object"}`) +} +func (f *fakeIdleTool) Execute(ctx context.Context, args json.RawMessage) (string, error) { + return f.exec(ctx, args) +} +func (f *fakeIdleTool) RequiresConfirmation(json.RawMessage) bool { return false } +func (f *fakeIdleTool) CallString(json.RawMessage) string { return f.name } + +// TestIdleWatchdogKillsHungToolBeforeAgent verifies the two-stage escalation +// end to end. A tool hung past the idle threshold stops masking idleness, so +// the idle event fires despite the in-flight call; on the tick that exceeds +// the kill threshold the watchdog kills the TOOL first (its cancellable +// context, registered by executor.ExecuteToolCalls, makes the call return a +// "tool cancelled" result) and the agent recovers — the next tool call runs on +// a still-alive parent context and the agent-level self-cancel never fires. +func TestIdleWatchdogKillsHungToolBeforeAgent(t *testing.T) { + o := newIdleTestOrchestrator(t, 60*time.Millisecond, 150*time.Millisecond, 50*time.Millisecond, []client.ChatMessage{ + {Role: "user", Content: client.TextContent("goal")}, + }) + + started := make(chan struct{}) + sawCancel := make(chan error, 1) + o.sess.Registry.Register(&fakeIdleTool{ + name: "hung_tool", + exec: func(ctx context.Context, args json.RawMessage) (string, error) { + close(started) + <-ctx.Done() + sawCancel <- ctx.Err() + return "", fmt.Errorf("tool cancelled: %v", ctx.Err()) + }, + }) + o.sess.Registry.Register(&fakeIdleTool{ + name: "quick_tool", + exec: func(ctx context.Context, args json.RawMessage) (string, error) { + if err := ctx.Err(); err != nil { + return "", fmt.Errorf("inherited a cancelled context: %v", err) + } + return "recovered", nil + }, + }) + + toolCalls := []client.ToolCall{ + {ID: "tc_1", Function: client.FunctionCall{Name: "hung_tool", Arguments: "{}"}}, + {ID: "tc_2", Function: client.FunctionCall{Name: "quick_tool", Arguments: "{}"}}, + } + + runCtx, cancelRun := context.WithCancel(context.Background()) + defer cancelRun() + o.startIdleWatchdog(runCtx, cancelRun) + + done := make(chan error, 1) + go func() { + done <- executor.ExecuteToolCalls(runCtx, o.sess, toolCalls, o.withActivityMiddleware()) + }() + + <-started + waitFor(t, 2*time.Second, func() bool { return o.inFlightTools.Load() == 1 }) + + // Nothing but the watchdog's stage-1 tool kill can unblock the hung call: + // once ExecuteToolCalls returns, the escalation has happened. + select { + case err := <-done: + if err != nil { + t.Fatalf("ExecuteToolCalls returned an error after the tool kill: %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("the watchdog did not kill the hung tool") + } + + if err := <-sawCancel; err != context.Canceled { + t.Fatalf("the hung tool's context error = %v, want context.Canceled", err) + } + + // The agent continued: the killed call recorded the cancellation result + // and the next call ran with a fresh context. + if len(o.sess.History) != 3 { + t.Fatalf("expected 3 history entries (user + 2 tool results), got %d", len(o.sess.History)) + } + if first := o.sess.History[1].Content.String(); first != "tool cancelled by the harness idle watchdog" { + t.Errorf("killed call result = %q, want the tool-cancellation message", first) + } + if second := o.sess.History[2].Content.String(); second != "recovered" { + t.Errorf("post-kill tool result = %q, want %q", second, "recovered") + } + + waitFor(t, 2*time.Second, func() bool { return o.inFlightTools.Load() == 0 }) + waitFor(t, 2*time.Second, func() bool { return o.oldestToolStartAt.Load() == 0 }) + + // Stage 2 must not fire: activity resumed with the next tool call, so this + // window ends well before the re-armed idle age could reach the kill + // threshold again (~killAfter after the recovery). Blocking on the event + // channel for the window lets any re-armed idle episode (harmless) pass. + windowEvents := collectIdleEvents(t, o, 80*time.Millisecond) + if runCtx.Err() != nil { + t.Errorf("agent run context was cancelled after the tool kill: %v", runCtx.Err()) + } + if reason := o.IdleKillReason(); reason != "" { + t.Errorf("idle kill reason recorded without an agent kill: %q", reason) + } + + // The idle event fired despite the in-flight (stalled) tool — the first + // episode fired before the kill; a re-armed follow-up episode may have + // fired since. Count what the window above collected plus whatever is + // still buffered. + idleEvents := len(windowEvents) +drainEvents: + for { + select { + case ev := <-o.eventCh: + if _, ok := ev.(common.SubagentIdleEvent); ok { + idleEvents++ + } + default: + break drainEvents + } + } + if idleEvents < 1 { + t.Fatal("expected at least one idle event despite the in-flight tool") + } +} + +// TestIdleKillSelfCancels verifies that sustained idle past the kill +// threshold cancels the orchestrator's own run context exactly once, with the +// probe lines recorded in the kill reason. +func TestIdleKillSelfCancels(t *testing.T) { + history := []client.ChatMessage{ + {Role: "user", Content: client.TextContent("goal")}, + {Role: "assistant", Content: client.TextContent("Reply 2")}, + } + o := newIdleTestOrchestrator(t, 40*time.Millisecond, 80*time.Millisecond, 10*time.Millisecond, history) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + o.startIdleWatchdog(ctx, cancel) + + select { + case <-ctx.Done(): + case <-time.After(2 * time.Second): + t.Fatal("run context was not cancelled by the idle watchdog") + } + if ctx.Err() != context.Canceled { + t.Fatalf("ctx.Err() = %v, want context.Canceled", ctx.Err()) + } + + reason := o.IdleKillReason() + if reason == "" { + t.Fatal("expected a recorded kill reason after the idle kill") + } + if !strings.Contains(reason, "assistant: Reply 2") { + t.Errorf("kill reason must include the probe lines, got %q", reason) + } + if !strings.Contains(reason, "idle for") { + t.Errorf("kill reason must report the idle duration, got %q", reason) + } + + // The kill happens on the same once-per-episode tick as the notification. + if events := collectIdleEvents(t, o, 50*time.Millisecond); len(events) != 1 { + t.Fatalf("expected exactly 1 idle event alongside the kill, got %d", len(events)) + } +} + +// TestMarkActivityThroughMiddleware verifies the activity middleware bumps +// lastActivity at entry and keeps inFlightTools > 0 for the whole execution. +func TestMarkActivityThroughMiddleware(t *testing.T) { + o := newIdleTestOrchestrator(t, 0, 0, 0, nil) + + before := time.Unix(0, o.lastActivity.Load()) + done := make(chan struct{}) + base := func(ctx context.Context, tc client.ToolCall) (string, error) { + time.Sleep(100 * time.Millisecond) + return "ok", nil + } + toolRunner := o.activityMiddleware()(base) + + go func() { + defer close(done) + _, _ = toolRunner(context.Background(), client.ToolCall{ + Type: "function", + Function: client.FunctionCall{Name: "fake_tool", Arguments: "{}"}, + }) + }() + + waitFor(t, 2*time.Second, func() bool { return o.inFlightTools.Load() > 0 }) + + if after := time.Unix(0, o.lastActivity.Load()); !after.After(before) { + t.Fatalf("lastActivity was not bumped by the middleware: before=%v after=%v", before, after) + } + // MarkActivity re-arms the idle notification for a fresh episode. + if o.idleNotified.Load() { + t.Fatal("MarkActivity must reset idleNotified") + } + if inFlight := o.inFlightTools.Load(); inFlight != 1 { + t.Fatalf("inFlightTools during execution = %d, want 1", inFlight) + } + + <-done + if inFlight := o.inFlightTools.Load(); inFlight != 0 { + t.Fatalf("inFlightTools after execution = %d, want 0", inFlight) + } +} + +// TestLastTranscriptLines covers the probe renderer: last-n window, role +// prefix, first-line-only content, tool-call fallback for content-less +// assistant messages, skipping fully empty messages, and length capping. +func TestLastTranscriptLines(t *testing.T) { + long := strings.Repeat("x", 300) + msgs := []client.ChatMessage{ + {Role: "user", Content: client.TextContent("dropped: outside the window")}, + {Role: "assistant", Content: client.TextContent("with tools")}, + {Role: "assistant", ToolCalls: []client.ToolCall{{Function: client.FunctionCall{Name: "bash"}}}}, + {Role: "tool", Content: client.TextContent("command output\nsecond line")}, + {Role: "assistant", Content: client.TextContent(long)}, + } + + lines := lastTranscriptLines(msgs, 3) + if len(lines) != 3 { + t.Fatalf("expected 3 probe lines, got %d: %v", len(lines), lines) + } + want := []string{ + "assistant: called bash", + "tool: command output", + } + for i, w := range want { + if lines[i] != w { + t.Errorf("probe[%d] = %q, want %q", i, lines[i], w) + } + } + if !strings.HasPrefix(lines[2], "assistant: ") || !strings.HasSuffix(lines[2], "...") { + t.Errorf("long line must be role-prefixed and truncated, got %q", lines[2]) + } + if len([]rune(lines[2])) > idleProbeLineMaxLen { + t.Errorf("probe line exceeds the cap: %d runes", len([]rune(lines[2]))) + } + if strings.Join(lines, "\n") == "" || strings.Contains(strings.Join(lines, "\n"), "dropped") { + t.Errorf("probe must only contain the last entries, got %v", lines) + } + + // A message with neither content nor tool calls contributes no line. + empty := []client.ChatMessage{{Role: "assistant"}} + if got := lastTranscriptLines(empty, 3); len(got) != 0 { + t.Errorf("expected no lines for empty messages, got %v", got) + } + if got := lastTranscriptLines(msgs, 0); got != nil { + t.Errorf("expected no lines for n=0, got %v", got) + } +} + +// TestIdleWatchdogDisabledWithoutThreshold pins the off switch: no policy, no +// watchdog activity, no events. +func TestIdleWatchdogDisabledWithoutThreshold(t *testing.T) { + o := newIdleTestOrchestrator(t, 0, 0, 10*time.Millisecond, []client.ChatMessage{ + {Role: "user", Content: client.TextContent("goal")}, + }) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + o.startIdleWatchdog(ctx, cancel) + + if events := collectIdleEvents(t, o, 100*time.Millisecond); len(events) != 0 { + t.Fatalf("watchdog must stay silent when idleTimeout is 0, got %v", events) + } +} diff --git a/internal/session/inflight.go b/internal/session/inflight.go new file mode 100644 index 00000000..4ff1a4b4 --- /dev/null +++ b/internal/session/inflight.go @@ -0,0 +1,33 @@ +package session + +import "context" + +// SetInFlightToolCancel registers the cancel function of the currently +// executing tool call, enabling the owning orchestrator (or the user) to +// cancel a hung tool without killing the whole agent. Nil clears it. +func (s *Session) SetInFlightToolCancel(cancel context.CancelFunc) { + if cancel == nil { + s.inFlightToolCancel.Store(nil) + return + } + s.inFlightToolCancel.Store(&cancel) +} + +// ClearInFlightToolCancel removes the registered in-flight tool cancel +// function. It must be called once a tool call has finished so a stale cancel +// is never fired against a later tool call. +func (s *Session) ClearInFlightToolCancel() { + s.inFlightToolCancel.Store(nil) +} + +// CancelInFlightTool cancels the in-flight tool call, if any. It returns +// false when no tool is in flight (no cancel function is registered); the +// caller can use that to distinguish "killed the tool" from "nothing to +// kill" when escalating a kill. +func (s *Session) CancelInFlightTool() bool { + if cancel := s.inFlightToolCancel.Load(); cancel != nil && *cancel != nil { + (*cancel)() + return true + } + return false +} diff --git a/internal/session/session.go b/internal/session/session.go index 7211f632..fb1c3635 100644 --- a/internal/session/session.go +++ b/internal/session/session.go @@ -10,6 +10,7 @@ import ( "path/filepath" "strings" "sync" + "sync/atomic" "time" ) @@ -25,6 +26,12 @@ type Session struct { subagentSeq int saveSubagentHistories *bool Registry *tool.Registry + + // inFlightToolCancel holds the cancel function of the tool call that is + // currently executing (nil when none), so the owning orchestrator — or the + // user — can kill a hung tool without killing the whole agent run. + // See inflight.go. + inFlightToolCancel atomic.Pointer[context.CancelFunc] } func New(c *client.Client, historyPath string, history []client.ChatMessage, systemPrompt string, useTools bool) *Session { diff --git a/internal/tool/implementations.go b/internal/tool/implementations.go index 876b0c87..edb867cb 100644 --- a/internal/tool/implementations.go +++ b/internal/tool/implementations.go @@ -304,6 +304,51 @@ const maxReadFileChars = 32768 // Maximum number of characters for shell output to prevent session poisoning const maxBashOutputChars = 32768 +// defaultShellTimeout bounds every shell invocation unless the caller overrides +// it with the per-call `timeout` parameter (or disables it with "0"). Tests +// override it via SetShellTimeout. +var defaultShellTimeout = 10 * time.Minute + +// SetShellTimeout overrides the global default shell timeout. +func SetShellTimeout(d time.Duration) { defaultShellTimeout = d } + +// resolveShellTimeout derives the effective execution context for a shell call +// from the optional per-call `timeout` parameter: +// +// - absent/empty → the global default timeout +// - "0" (or negative) → unlimited (no deadline; group-kill + WaitDelay still +// apply on explicit context cancellation) +// - a valid duration → a deadline of that duration +// - an invalid string → an error +// +// The returned duration is the effective bound in force (0 means unlimited) so +// callers can report accurate timeouts. The returned cancel func is nil when no +// deadline was installed. +func resolveShellTimeout(ctx context.Context, raw string) (context.Context, context.CancelFunc, time.Duration, error) { + raw = strings.TrimSpace(raw) + + if raw == "" { + if defaultShellTimeout > 0 { + execCtx, cancel := context.WithTimeout(ctx, defaultShellTimeout) + return execCtx, cancel, defaultShellTimeout, nil + } + return ctx, nil, 0, nil + } + + d, err := time.ParseDuration(raw) + if err != nil { + return nil, nil, 0, fmt.Errorf("invalid timeout %q — use a duration like 30m, 2h, or 0 for unlimited", raw) + } + if d <= 0 { + // "0" (or negative) means unlimited: run without a deadline. Explicit + // cancellation still kills the process group (see newShellCommand). + return ctx, nil, 0, nil + } + + execCtx, cancel := context.WithTimeout(ctx, d) + return execCtx, cancel, d, nil +} + // ShellTool executes host-native shell commands with security restrictions. type ShellTool struct{} @@ -323,7 +368,8 @@ func (t ShellTool) Parameters() json.RawMessage { "type": "object", "properties": { "command": { "type": "string", "description": "The full %s command to execute." }, - "cwd": { "type": "string", "description": "Working directory for execution. Use this instead of 'cd' commands to change directories." } + "cwd": { "type": "string", "description": "Working directory for execution. Use this instead of 'cd' commands to change directories." }, + "timeout": { "type": "string", "description": "Optional per-call time bound, e.g. 30m or 2h. 0 means unlimited. Defaults to the configured global timeout." } }, "required": ["command"] }`, shellDisplayName())) @@ -332,11 +378,22 @@ func (t ShellTool) Execute(ctx context.Context, args json.RawMessage) (string, e var params struct { Command string `json:"command"` Cwd string `json:"cwd"` + Timeout string `json:"timeout"` } if err := json.Unmarshal(args, ¶ms); err != nil { return "", err } + // Resolve the effective time bound for this call. An invalid timeout is a + // tool-level error result: nothing runs and no error sandwich is attached. + execCtx, cancel, effectiveTimeout, timeoutErr := resolveShellTimeout(ctx, params.Timeout) + if timeoutErr != nil { + return fmt.Sprintf("Error: %v", timeoutErr), nil + } + if cancel != nil { + defer cancel() + } + // Validate command before any execution if err := t.ValidateBashCommand(params.Command, params.Cwd); err != nil { return "", t.WrapError(ctx, err) @@ -366,11 +423,18 @@ func (t ShellTool) Execute(ctx context.Context, args json.RawMessage) (string, e } // Execute command using a platform-specific shell wrapper. - cmd := newShellCommand(ctx, params.Command) + cmd := newShellCommand(execCtx, params.Command) cmd.Dir = params.Cwd output, err := cmd.CombinedOutput() + // The deadline fired and the process group was killed: report the timeout + // plus whatever output made it out before the kill. Explicit cancellation + // (context.Canceled) keeps flowing through the normal error path below. + if err != nil && execCtx.Err() == context.DeadlineExceeded { + return "", fmt.Errorf("command timed out after %s and was killed (partial output):\n%s", effectiveTimeout, string(output)) + } + // If sqz is available, compress the output if IsSqzAvailable() && len(output) > 0 { compressed, sqzErr := CompressWithSqz(ctx, output, params.Command) @@ -411,7 +475,7 @@ func (t ShellTool) Execute(ctx context.Context, args json.RawMessage) (string, e if orchestratorID := common.GetOrchestratorID(ctx); strings.Contains(strings.ToLower(orchestratorID), "coder") { sandwich = "\n\n=========================================\nSYSTEM DIRECTIVE:\nYou just encountered an error. If fixing this requires modifying components or architecture you were not explicitly instructed to edit, YOU MUST ABORT AND RETURN TO THE MAIN AGENT.\n=========================================" } - + if exitErr, ok := err.(*exec.ExitError); ok { return fmt.Sprintf("Command failed with exit code %d\n%s%s", exitErr.ExitCode(), finalOutput, sandwich), nil } diff --git a/internal/tool/shell_command_unix.go b/internal/tool/shell_command_unix.go index 2dcfc0f6..f8a14ee7 100644 --- a/internal/tool/shell_command_unix.go +++ b/internal/tool/shell_command_unix.go @@ -6,6 +6,8 @@ import ( "context" "os/exec" "sync" + "syscall" + "time" ) var ( @@ -29,5 +31,28 @@ func getUnixShellPath() string { } func newShellCommand(ctx context.Context, command string) *exec.Cmd { - return exec.CommandContext(ctx, getUnixShellPath(), "-c", command) + cmd := exec.CommandContext(ctx, getUnixShellPath(), "-c", command) + + // Run the shell in its own process group and kill the whole group on + // cancellation. exec.CommandContext only signals the direct child, so a + // grandchild that inherited stdout/stderr (the pipes Wait reads) would + // survive the shell and block Wait forever; a group-wide SIGKILL reaps the + // entire tree and WaitDelay (below) caps any pipe still held afterwards. + cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true} + cmd.Cancel = func() error { + if cmd.Process == nil { + return nil + } + // Negative PID targets the process group; ESRCH means the group is + // already gone, which is success for a kill. + if err := syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL); err != nil && err != syscall.ESRCH { + return err + } + return nil + } + // Give up waiting on pipe-holding descendants 5s after the group kill so + // Wait (and therefore CombinedOutput) always returns. + cmd.WaitDelay = 5 * time.Second + + return cmd } diff --git a/internal/tool/shell_command_windows.go b/internal/tool/shell_command_windows.go index 7590c6cc..98d70a2e 100644 --- a/internal/tool/shell_command_windows.go +++ b/internal/tool/shell_command_windows.go @@ -6,7 +6,9 @@ import ( "context" "encoding/base64" "os/exec" + "strconv" "sync" + "time" "unicode/utf16" ) @@ -43,9 +45,27 @@ func encodePSCommand(command string) string { func newShellCommand(ctx context.Context, command string) *exec.Cmd { shell := getWindowsShellPath() encoded := encodePSCommand(command) - return exec.CommandContext( + cmd := exec.CommandContext( ctx, shell, "-NoProfile", "-NonInteractive", "-ExecutionPolicy", "Bypass", "-EncodedCommand", encoded, ) + + // exec.CommandContext kills only the direct PowerShell child; descendants + // that inherited the output pipes would keep Wait blocked forever after the + // shell itself is gone. taskkill /T walks the process tree instead. + cmd.Cancel = func() error { + if cmd.Process == nil { + return nil + } + // Best-effort tree kill: taskkill errors are ignored because WaitDelay + // (below) still guarantees Wait returns. + _ = exec.Command("taskkill", "/T", "/F", "/PID", strconv.Itoa(cmd.Process.Pid)).Run() + return nil + } + // Grace period for descendants that keep the output pipes open after the + // kill, so Wait (and therefore CombinedOutput) always returns. + cmd.WaitDelay = 5 * time.Second + + return cmd } diff --git a/internal/tool/shell_timeout_test.go b/internal/tool/shell_timeout_test.go new file mode 100644 index 00000000..ce2f85cb --- /dev/null +++ b/internal/tool/shell_timeout_test.go @@ -0,0 +1,226 @@ +//go:build !windows + +package tool + +import ( + "context" + "encoding/json" + "strings" + "testing" + "time" +) + +// setShellTimeoutForTest overrides the global default shell timeout and +// restores the previous value when the test finishes. +func setShellTimeoutForTest(t *testing.T, d time.Duration) { + t.Helper() + prev := defaultShellTimeout + SetShellTimeout(d) + t.Cleanup(func() { SetShellTimeout(prev) }) +} + +// TestShellTool_TimeoutKillsHangingCommand verifies the global default timeout: +// a hanging command is killed and the error reports the timeout plus the +// partial output produced before the kill. +func TestShellTool_TimeoutKillsHangingCommand(t *testing.T) { + setShellTimeoutForTest(t, 2*time.Second) + + tool := ShellTool{} + args := json.RawMessage(`{"command": "echo start; sleep 300"}`) + + start := time.Now() + _, err := tool.Execute(approvedContext(), args) + elapsed := time.Since(start) + + if err == nil { + t.Fatal("expected timeout error, got nil") + } + if !strings.Contains(err.Error(), "timed out") { + t.Fatalf("expected error to mention the timeout, got %q", err.Error()) + } + if !strings.Contains(err.Error(), "start") { + t.Fatalf("expected partial output 'start' in the timeout error, got %q", err.Error()) + } + if elapsed >= 5*time.Second { + t.Fatalf("command should have been killed after ~2s, took %s", elapsed) + } +} + +// TestShellTool_CancelReturnsDespiteGrandchildHoldingPipes verifies that an +// explicit context cancellation returns promptly even when the command spawns +// descendants: the process-group kill reaps the tree and WaitDelay caps any +// pipe-holding straggler. +func TestShellTool_CancelReturnsDespiteGrandchildHoldingPipes(t *testing.T) { + setShellTimeoutForTest(t, 0) // unlimited; the bound comes from ctx cancellation + + tool := ShellTool{} + args := json.RawMessage(`{"command": "echo start; (sleep 300 &) ; sleep 300"}`) + + ctx, cancel := context.WithCancel(approvedContext()) + defer cancel() + + type outcome struct { + result string + err error + } + outcomeCh := make(chan outcome, 1) + go func() { + result, err := tool.Execute(ctx, args) + outcomeCh <- outcome{result: result, err: err} + }() + + time.Sleep(200 * time.Millisecond) + cancelAt := time.Now() + cancel() + + select { + case <-outcomeCh: + case <-time.After(10 * time.Second): + // Fail-safe: WaitDelay is 5s, so a hang here means the group kill did + // not fire and pipes are holding Wait open. + t.Fatal("Execute did not return within 10s of cancellation") + } + if elapsed := time.Since(cancelAt); elapsed >= 9*time.Second { + t.Fatalf("Execute took %s to return after cancellation, want < 9s", elapsed) + } +} + +// TestShellTool_PerCallTimeoutOverridesGlobal verifies the per-call timeout +// parameter beats the (longer) global default. +func TestShellTool_PerCallTimeoutOverridesGlobal(t *testing.T) { + setShellTimeoutForTest(t, 10*time.Minute) + + tool := ShellTool{} + args := json.RawMessage(`{"command": "echo start; sleep 300", "timeout": "1s"}`) + + start := time.Now() + _, err := tool.Execute(approvedContext(), args) + elapsed := time.Since(start) + + if err == nil { + t.Fatal("expected timeout error, got nil") + } + if !strings.Contains(err.Error(), "timed out") { + t.Fatalf("expected error to mention the timeout, got %q", err.Error()) + } + if !strings.Contains(err.Error(), "start") { + t.Fatalf("expected partial output 'start' in the timeout error, got %q", err.Error()) + } + if elapsed >= 5*time.Second { + t.Fatalf("command should have been killed after ~1s, took %s", elapsed) + } +} + +// TestShellTool_PerCallTimeoutUnlimited verifies that "0" parses as unlimited +// without hanging a fast command. +func TestShellTool_PerCallTimeoutUnlimited(t *testing.T) { + setShellTimeoutForTest(t, 10*time.Minute) + + tool := ShellTool{} + args := json.RawMessage(`{"command": "echo ok", "timeout": "0"}`) + + result, err := tool.Execute(approvedContext(), args) + if err != nil { + t.Fatalf("expected success with unlimited timeout, got %v", err) + } + if !strings.Contains(result, "ok") { + t.Fatalf("expected command output 'ok', got %q", result) + } +} + +// TestShellTool_InvalidTimeoutIsErrorResult verifies that an unparsable timeout +// is reported as an error RESULT (nil Go error) and that the command never runs. +func TestShellTool_InvalidTimeoutIsErrorResult(t *testing.T) { + setShellTimeoutForTest(t, 10*time.Minute) + + tool := ShellTool{} + args := json.RawMessage(`{"command": "echo should-not-run", "timeout": "banana"}`) + + result, err := tool.Execute(approvedContext(), args) + if err != nil { + t.Fatalf("expected an error result with nil Go error, got %v", err) + } + if !strings.Contains(result, "invalid timeout") { + t.Fatalf("expected result to report the invalid timeout, got %q", result) + } + if strings.Contains(result, "should-not-run") { + t.Fatalf("command must not run when the timeout is invalid, got %q", result) + } +} + +// TestResolveShellTimeout covers the pure timeout resolution function directly. +func TestResolveShellTimeout(t *testing.T) { + tests := []struct { + name string + raw string + wantDuration time.Duration + wantDeadline bool + wantCancelFunc bool + wantErr bool + }{ + { + name: "empty uses global default", + raw: "", + wantDuration: 10 * time.Minute, + wantDeadline: true, + wantCancelFunc: true, + }, + { + name: "zero is unlimited", + raw: "0", + wantDuration: 0, + wantDeadline: false, + }, + { + name: "negative is unlimited", + raw: "-5m", + wantDuration: 0, + wantDeadline: false, + }, + { + name: "valid duration is honored", + raw: "90m", + wantDuration: 90 * time.Minute, + wantDeadline: true, + wantCancelFunc: true, + }, + { + name: "invalid string errors", + raw: "banana", + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + setShellTimeoutForTest(t, 10*time.Minute) + + ctx, cancel, gotDuration, err := resolveShellTimeout(context.Background(), tt.raw) + if tt.wantErr { + if err == nil { + t.Fatalf("expected an error for raw=%q, got nil", tt.raw) + } + if !strings.Contains(err.Error(), "invalid timeout") { + t.Fatalf("expected error to mention 'invalid timeout', got %q", err.Error()) + } + return + } + if err != nil { + t.Fatalf("unexpected error for raw=%q: %v", tt.raw, err) + } + if gotDuration != tt.wantDuration { + t.Fatalf("resolved duration = %s, want %s", gotDuration, tt.wantDuration) + } + _, hasDeadline := ctx.Deadline() + if hasDeadline != tt.wantDeadline { + t.Fatalf("ctx hasDeadline = %v, want %v", hasDeadline, tt.wantDeadline) + } + if (cancel != nil) != tt.wantCancelFunc { + t.Fatalf("cancel func presence = %v, want %v", cancel != nil, tt.wantCancelFunc) + } + if cancel != nil { + cancel() + } + }) + } +} diff --git a/internal/tool/subagent.go b/internal/tool/subagent.go index 244e150c..5f74482e 100644 --- a/internal/tool/subagent.go +++ b/internal/tool/subagent.go @@ -5,11 +5,17 @@ import ( "encoding/json" "fmt" "strings" + "time" "late/internal/assets" ) -type SubagentRunner func(ctx context.Context, goal string, ctxFiles []string, agentType string) (string, error) +// SubagentRunner executes one subagent run. timeoutOverride carries the +// per-spawn wall-clock budget parsed from the spawn_subagent "timeout" +// argument: nil = no override (the global --subagent-timeout/config value +// applies); a non-positive value = unlimited (no budget for this run); +// a positive value = a per-spawn budget overriding the global one. +type SubagentRunner func(ctx context.Context, goal string, ctxFiles []string, agentType string, timeoutOverride *time.Duration) (string, error) type SpawnSubagentTool struct { Runner SubagentRunner @@ -44,6 +50,10 @@ func (t SpawnSubagentTool) Parameters() json.RawMessage { "type": "string", "enum": [%s], "description": "The type of subagent to spawn. %s" + }, + "timeout": { + "type": "string", + "description": "Optional wall-clock budget for this subagent run, e.g. \"45m\", \"2h\"; \"0\" = unlimited; omitted = the global --subagent-timeout/config value" } }, "required": ["goal", "agent_type"] @@ -61,12 +71,40 @@ func (t SpawnSubagentTool) Execute(ctx context.Context, args json.RawMessage) (s Goal string `json:"goal"` CtxFiles []string `json:"ctx_files"` AgentType string `json:"agent_type"` + Timeout string `json:"timeout"` } if err := json.Unmarshal(args, ¶ms); err != nil { return "", fmt.Errorf("failed to parse arguments: %v", err) } - return t.Runner(ctx, params.Goal, params.CtxFiles, params.AgentType) + timeoutOverride, err := parseSubagentTimeout(params.Timeout) + if err != nil { + // Surface the failure as an error RESULT (nil Go error) so the model + // can read the hint and retry with a valid duration. + return fmt.Sprintf("Error: invalid subagent timeout %q — use a duration like 45m, 2h, or 0 for unlimited", params.Timeout), nil + } + + return t.Runner(ctx, params.Goal, params.CtxFiles, params.AgentType, timeoutOverride) +} + +// parseSubagentTimeout parses the optional per-spawn "timeout" argument. +// Empty (absent) → nil override: the global budget applies. +// "0" or a negative duration → pointer to 0 (explicit unlimited). +// A positive duration → pointer to that budget. +// Anything else → a parse error for the caller to surface as an error result. +func parseSubagentTimeout(raw string) (*time.Duration, error) { + if raw == "" { + return nil, nil + } + parsed, err := time.ParseDuration(raw) + if err != nil { + return nil, err + } + if parsed <= 0 { + unlimited := time.Duration(0) + return &unlimited, nil + } + return &parsed, nil } func (t SpawnSubagentTool) RequiresConfirmation(args json.RawMessage) bool { return false } diff --git a/internal/tool/subagent_test.go b/internal/tool/subagent_test.go new file mode 100644 index 00000000..bf00c3c6 --- /dev/null +++ b/internal/tool/subagent_test.go @@ -0,0 +1,130 @@ +package tool + +import ( + "context" + "encoding/json" + "strings" + "testing" + "time" +) + +// TestSpawnSubagentTool_TimeoutParsing covers the per-spawn "timeout" +// argument: absent/empty → no override; a valid duration → parsed override; +// "0"/negative → explicit unlimited (pointer to 0); an invalid value → an +// error RESULT (nil Go error) with the runner never invoked, so the model +// can retry with a valid duration. +func TestSpawnSubagentTool_TimeoutParsing(t *testing.T) { + tests := []struct { + name string + args string + wantErrStr bool // expect an error-result string; runner not called + wantOverride *time.Duration // expected override passed to the runner + }{ + { + name: "absent timeout means no override", + args: `{"goal":"g","agent_type":"coder"}`, + wantOverride: nil, + }, + { + name: "empty timeout means no override", + args: `{"goal":"g","agent_type":"coder","timeout":""}`, + wantOverride: nil, + }, + { + name: "valid duration is passed through", + args: `{"goal":"g","agent_type":"coder","timeout":"45m"}`, + wantOverride: subagentTimeoutPtr(45 * time.Minute), + }, + { + name: "two hours is passed through", + args: `{"goal":"g","agent_type":"coder","timeout":"2h"}`, + wantOverride: subagentTimeoutPtr(2 * time.Hour), + }, + { + name: "zero means explicit unlimited", + args: `{"goal":"g","agent_type":"coder","timeout":"0"}`, + wantOverride: subagentTimeoutPtr(0), + }, + { + name: "negative means explicit unlimited normalized to 0", + args: `{"goal":"g","agent_type":"coder","timeout":"-5m"}`, + wantOverride: subagentTimeoutPtr(0), + }, + { + name: "garbage duration is an error result", + args: `{"goal":"g","agent_type":"coder","timeout":"banana"}`, + wantErrStr: true, + }, + { + name: "missing unit is an error result", + args: `{"goal":"g","agent_type":"coder","timeout":"5"}`, + wantErrStr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + runnerCalled := false + var gotOverride *time.Duration + spawnTool := SpawnSubagentTool{ + Runner: func(ctx context.Context, goal string, ctxFiles []string, agentType string, timeoutOverride *time.Duration) (string, error) { + runnerCalled = true + gotOverride = timeoutOverride + if goal != "g" || agentType != "coder" { + t.Errorf("runner got goal=%q agentType=%q, want goal=%q agentType=%q", goal, agentType, "g", "coder") + } + return "ok", nil + }, + } + + result, err := spawnTool.Execute(context.Background(), json.RawMessage(tt.args)) + if err != nil { + t.Fatalf("Execute() Go error = %v, want nil", err) + } + + if tt.wantErrStr { + if runnerCalled { + t.Fatal("runner must not be invoked for an invalid timeout") + } + if !strings.Contains(result, `invalid subagent timeout`) || !strings.Contains(result, "use a duration like 45m, 2h") { + t.Fatalf("error result = %q, want the invalid-timeout retry hint", result) + } + return + } + + if !runnerCalled { + t.Fatal("runner was not invoked") + } + if tt.wantOverride == nil { + if gotOverride != nil { + t.Fatalf("override = %v, want nil", *gotOverride) + } + return + } + if gotOverride == nil { + t.Fatalf("override = nil, want %v", *tt.wantOverride) + } + if *gotOverride != *tt.wantOverride { + t.Fatalf("override = %v, want %v", *gotOverride, *tt.wantOverride) + } + }) + } +} + +// TestSpawnSubagentTool_ParametersDocumentTimeout guards the JSON schema: +// the optional timeout property and its budget semantics must stay advertised +// to the model. +func TestSpawnSubagentTool_ParametersDocumentTimeout(t *testing.T) { + schema := string(SpawnSubagentTool{Runner: nil}.Parameters()) + if !strings.Contains(schema, `"timeout"`) { + t.Fatal("parameters schema does not advertise the timeout property") + } + if !strings.Contains(schema, "unlimited") { + t.Fatal("timeout schema description does not document the unlimited semantics") + } + if !strings.Contains(schema, "--subagent-timeout/config value") { + t.Fatal("timeout schema description does not document the omitted-means-global semantics") + } +} + +func subagentTimeoutPtr(d time.Duration) *time.Duration { return &d } diff --git a/internal/tui/update.go b/internal/tui/update.go index 69789299..c1724ec2 100644 --- a/internal/tui/update.go +++ b/internal/tui/update.go @@ -1586,6 +1586,20 @@ func (m Model) updateChat(msg tea.Msg) (Model, tea.Cmd) { if event.ID == m.Focused.ID() { m.updateViewport() } + case common.SubagentIdleEvent: + // Status line only: the agent is still formally running, so the + // State must not change. The event fires once per idle episode. + first := "" + if len(event.Probe) > 0 { + first = event.Probe[0] + } + if r := []rune(first); len(r) > 80 { + first = string(r[:77]) + "..." + } + s.StatusText = fmt.Sprintf("subagent idle for %s — last: %s", event.IdleFor.Truncate(time.Second), first) + if event.ID == m.Focused.ID() { + m.updateViewport() + } } case ConfirmRequestMsg: