From 791a5a5e3ecaa46890a30a6d12c359c5d816ecbb Mon Sep 17 00:00:00 2001 From: Lakshman Patel Date: Thu, 6 Aug 2026 08:18:27 +0530 Subject: [PATCH] feat(tasks): background TaskRunner executor + expanded task states MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes the last KiroCrew gap: the TaskStore was store-only — nothing consumed its retry/replan/checkpoint machinery. Add the execution half: - TaskRunner (internal/tool/task_executor.go): a background executor loop that picks up ready work (pending, no open blockers), runs each task through a host Execute callback, applies the store's retry budget with exponential backoff via MarkFailed, replans tasks that exhaust their budget (OnReplan + Requeue, bounded by MaxReplans), enforces a MaxTotalTasks watchdog, and cancels in-flight tasks on stop. - New task states (KiroCrew parity): reviewing, skipped, cancelled, with TaskStore.Skip/Cancel and Schedule() validation extended. - TaskRunTool + ToolContext.TaskExecutor: the agent (or host) can invoke TaskRun to drive ready tasks; the session arms it with an agent-spawn based executor (subject/description/checkpoint become the sub-agent prompt). TaskUpdateTool accepts the new states and resets the retry budget on any explicit exit from failed. --- cmd/chat_tools.go | 1 + internal/engine/session.go | 45 +++ internal/engine/tool_service.go | 2 + internal/tool/task_create.go | 9 +- internal/tool/task_executor.go | 485 ++++++++++++++++++++++++++++ internal/tool/task_executor_test.go | 402 +++++++++++++++++++++++ internal/tool/task_run_tool.go | 83 +++++ internal/tool/task_run_tool_test.go | 69 ++++ internal/tool/task_runner.go | 58 ++++ internal/tool/task_schedule.go | 3 +- internal/tool/tool.go | 3 + 11 files changed, 1156 insertions(+), 4 deletions(-) create mode 100644 internal/tool/task_executor.go create mode 100644 internal/tool/task_executor_test.go create mode 100644 internal/tool/task_run_tool.go create mode 100644 internal/tool/task_run_tool_test.go diff --git a/cmd/chat_tools.go b/cmd/chat_tools.go index dcfa3bb7..b0157819 100644 --- a/cmd/chat_tools.go +++ b/cmd/chat_tools.go @@ -96,6 +96,7 @@ func optionalTools() []tool.Tool { tool.TaskGetTool{}, tool.TaskListTool{}, tool.TaskUpdateTool{}, + tool.TaskRunTool{}, tool.SleepTool{}, tool.CronCreateTool{}, tool.CronDeleteTool{}, diff --git a/internal/engine/session.go b/internal/engine/session.go index 50b252ce..ee947623 100644 --- a/internal/engine/session.go +++ b/internal/engine/session.go @@ -2,6 +2,7 @@ package engine import ( "context" + "encoding/json" "fmt" "log/slog" "os" @@ -184,6 +185,7 @@ func NewSessionWithClient(chat ChatClient, provider, model, systemPrompt string, recordVerification: s.recordVerificationObservation, lifecycle: s.life, appendSystem: s.AppendSystemContext, + taskExec: s.taskExecFromAgentSpawn(), }) s.refreshContextWindowCache() s.life.SetAgentsAccumulator(agentsAccum) @@ -209,6 +211,49 @@ func (s *Session) ReattachTransport(chat ChatClient, provider string, deployment _ = deploymentRouting } +// taskExecFromAgentSpawn returns the TaskRun executor: it spawns a general +// sub-agent to perform a stored task, feeding the task's description, active +// form, and checkpoint as the prompt. Nil agent-spawn capability disables the +// executor (the TaskRun tool then reports that no executor is configured). +func (s *Session) taskExecFromAgentSpawn() tool.TaskExecutorFunc { + return func(ctx context.Context, t *tool.Task) (string, error) { + if s == nil || s.tools == nil || s.tools.AgentSpawnFn() == nil { + return "", fmt.Errorf("task execution unavailable: no agent spawn capability") + } + prompt := "Execute the following task and report results.\n\nSubject: " + t.Subject + + "\n\nDescription: " + t.Description + if t.ActiveForm != "" { + prompt += "\n\n(You are working on: " + t.ActiveForm + ")" + } + if len(t.Checkpoint) > 0 { + b, _ := json.Marshal(t.Checkpoint) + prompt += "\n\nPrior progress (checkpoint): " + string(b) + } + res, err := s.tools.AgentSpawnFn()(ctx, agentcontracts.SpawnRequest{ + Prompt: prompt, + Description: "Execute task " + t.ID, + SubagentType: "general", + }) + if err != nil { + return "", err + } + if res.Status == agentcontracts.StatusFailed { + if res.Output != "" { + return "", fmt.Errorf("%s", res.Output) + } + return "", fmt.Errorf("task agent reported failure") + } + out := res.Output + if res.Summary != "" { + if out != "" { + out += "\n" + } + out += res.Summary + } + return out, nil + } +} + // SubSession clones transport and routing mode for explore/general sub-agents. func (s *Session) SubSession(model, systemPrompt string, registry *tool.Registry) *Session { if registry == nil { diff --git a/internal/engine/tool_service.go b/internal/engine/tool_service.go index 4e10a018..5ec166a9 100644 --- a/internal/engine/tool_service.go +++ b/internal/engine/tool_service.go @@ -70,6 +70,7 @@ type toolExecutionDeps struct { recordVerification func(types.ToolCall, string, bool) lifecycle *LifecycleService appendSystem func(string) + taskExec tool.TaskExecutorFunc } // NewToolService constructs a ToolService with the given registry. @@ -395,6 +396,7 @@ func (s *ToolService) ExecuteOne(ctx context.Context, tc types.ToolCall, overrid AvailableTools: available, Registry: s.registry, AutoCommit: s.AutoCommit(), + TaskExecutor: s.deps.taskExec, }) // Bridge session sandbox policy onto the context so Bash/PowerShell // WrapCommand actually applies. Path guards already read ToolContext.SandboxMode; diff --git a/internal/tool/task_create.go b/internal/tool/task_create.go index d6818568..3c4bbd9b 100644 --- a/internal/tool/task_create.go +++ b/internal/tool/task_create.go @@ -15,8 +15,11 @@ type TaskStatus string const ( TaskStatusPending TaskStatus = "pending" TaskStatusInProgress TaskStatus = "in_progress" + TaskStatusReviewing TaskStatus = "reviewing" TaskStatusCompleted TaskStatus = "completed" TaskStatusFailed TaskStatus = "failed" + TaskStatusSkipped TaskStatus = "skipped" + TaskStatusCancelled TaskStatus = "cancelled" ) // DefaultMaxAttempts is the retry budget used when a task does not declare @@ -399,7 +402,7 @@ func (TaskUpdateTool) Parameters() map[string]interface{} { "type": "object", "properties": map[string]interface{}{ "taskId": map[string]interface{}{"type": "string", "description": "The ID of the task to update"}, - "status": map[string]interface{}{"type": "string", "enum": []string{"pending", "in_progress", "completed", "failed"}, "description": "New task status"}, + "status": map[string]interface{}{"type": "string", "enum": []string{"pending", "in_progress", "reviewing", "completed", "failed", "skipped", "cancelled"}, "description": "New task status"}, "owner": map[string]interface{}{"type": "string", "description": "Agent name to assign"}, "dependencies": map[string]interface{}{"type": "array", "items": map[string]interface{}{"type": "object", "properties": map[string]interface{}{"targetId": map[string]interface{}{"type": "string"}, "type": map[string]interface{}{"type": "string", "enum": []string{"blocks", "related", "parent-child"}}}}, "description": "Replace dependencies"}, }, @@ -422,9 +425,9 @@ func (TaskUpdateTool) Execute(_ context.Context, input json.RawMessage) (string, } ok := globalTaskStore.Update(p.TaskID, func(t *Task) { if p.Status != "" { - // Moving a failed task back to pending is an explicit replan + // Any explicit status change out of failed is a replan or terminal // signal: reset the retry budget and error so it starts fresh. - if TaskStatus(p.Status) == TaskStatusPending && t.Status == TaskStatusFailed { + if t.Status == TaskStatusFailed && TaskStatus(p.Status) != TaskStatusFailed { t.Attempts = 0 t.LastError = "" } diff --git a/internal/tool/task_executor.go b/internal/tool/task_executor.go new file mode 100644 index 00000000..55ab5a0f --- /dev/null +++ b/internal/tool/task_executor.go @@ -0,0 +1,485 @@ +package tool + +import ( + "context" + "fmt" + "sync" + "time" +) + +// TaskExecutorFunc runs a single task and returns its result output. The +// context is cancelled when the run is stopped or the per-task timeout elapses; +// implementations should respect it. +type TaskExecutorFunc func(ctx context.Context, task *Task) (string, error) + +// TaskRunnerEvent is emitted through OnProgress when a task's lifecycle state +// changes. Events: "started", "completed", "retrying", "replanned", "failed", +// "skipped", "cancelled". +type TaskRunnerEvent string + +const ( + EventStarted TaskRunnerEvent = "started" + EventCompleted TaskRunnerEvent = "completed" + EventRetrying TaskRunnerEvent = "retrying" + EventReplanned TaskRunnerEvent = "replanned" + EventFailed TaskRunnerEvent = "failed" + EventSkipped TaskRunnerEvent = "skipped" + EventCancelled TaskRunnerEvent = "cancelled" +) + +// TaskRunnerOptions configures the TaskRunner. Zero values fall back to sane +// defaults; only Store and Execute are required. +type TaskRunnerOptions struct { + // Store is the task store the runner drives. Required. + Store *TaskStore + // Execute runs one task and returns its output. Required. + Execute TaskExecutorFunc + // OnProgress is called (outside locks) on each lifecycle transition. + OnProgress func(task *Task, event TaskRunnerEvent) + // OnReplan is invoked when a task parks failed (retry budget exhausted). + // Return true to requeue the task with a fresh retry budget (the + // "replan remaining work" behavior). Bounded by MaxReplans per task. + OnReplan func(ctx context.Context, task *Task) (bool, error) + // MaxTotalTasks is the watchdog cap on distinct tasks that may reach a + // terminal state in one run (default 50). Guards against unbounded work. + MaxTotalTasks int + // MaxReplans caps replans per task (default 2). + MaxReplans int + // PollInterval is the idle poll period (default 200ms). + PollInterval time.Duration + // Concurrency caps parallel task executions (default 1). + Concurrency int + // Backoff returns the wait before retrying attempt n (default quadratic, + // capped at 30s). + Backoff func(attempt int) time.Duration + // DefaultTimeout bounds each task execution (0 = no timeout). + DefaultTimeout time.Duration +} + +// TaskRunnerStats is a point-in-time snapshot of runner activity. +type TaskRunnerStats struct { + Running bool + Executed int + Completed int + Failed int + Skipped int + Cancelled int + Replanned int +} + +const ( + defaultMaxTotalTasks = 50 + defaultMaxReplans = 2 + defaultPollInterval = 200 * time.Millisecond + defaultConcurrency = 1 + defaultBackoffCap = 30 * time.Second +) + +// TaskRunner is a background executor that drives a TaskStore: it repeatedly +// picks up ready work (pending tasks with no open blockers), executes each task +// through Execute, applies the store's retry budget on failure (with backoff), +// replans tasks that exhaust their budget, and stops under a watchdog cap or on +// cancellation. It is the execution half of the store-only TaskStore. +type TaskRunner struct { + mu sync.Mutex + store *TaskStore + opts TaskRunnerOptions + backoff func(attempt int) time.Duration + + started bool + finished bool + cancel context.CancelFunc + doneCh chan struct{} + + replansByTask map[string]int + + executed int + completed int + failed int + skipped int + cancelled int + replanned int +} + +// NewTaskRunner validates and fills options with defaults. +func NewTaskRunner(opts TaskRunnerOptions) *TaskRunner { + if opts.MaxTotalTasks <= 0 { + opts.MaxTotalTasks = defaultMaxTotalTasks + } + if opts.MaxReplans < 0 { + opts.MaxReplans = 0 + } + if opts.MaxReplans == 0 { + opts.MaxReplans = defaultMaxReplans + } + if opts.PollInterval <= 0 { + opts.PollInterval = defaultPollInterval + } + if opts.Concurrency <= 0 { + opts.Concurrency = defaultConcurrency + } + if opts.Backoff == nil { + opts.Backoff = defaultBackoff + } + return &TaskRunner{ + store: opts.Store, + opts: opts, + backoff: opts.Backoff, + replansByTask: make(map[string]int), + } +} + +func defaultBackoff(attempt int) time.Duration { + if attempt <= 0 { + return 0 + } + d := time.Duration(attempt) * time.Duration(attempt) * 500 * time.Millisecond + if d > defaultBackoffCap { + d = defaultBackoffCap + } + return d +} + +// Run drives the store until quiescence (no ready work and nothing in flight), +// the watchdog cap, or ctx cancellation — whichever comes first. It is safe to +// call directly (blocking) or via Start. +func (r *TaskRunner) Run(ctx context.Context) error { + if r == nil || r.store == nil { + return fmt.Errorf("task runner: store is required") + } + if r.opts.Execute == nil { + return fmt.Errorf("task runner: execute function is required") + } + + ticker := time.NewTicker(r.opts.PollInterval) + defer ticker.Stop() + + for { + select { + case <-ctx.Done(): + r.cancelActive("cancelled") + return nil + default: + } + + if r.terminalCount() >= r.opts.MaxTotalTasks { + return nil + } + + ready := r.store.GetReadyWork() + if len(ready) == 0 { + if r.quiescent() { + return nil + } + select { + case <-ctx.Done(): + r.cancelActive("cancelled") + return nil + case <-ticker.C: + } + continue + } + + due := r.dueTasks(ready, time.Now()) + if len(due) == 0 { + select { + case <-ctx.Done(): + r.cancelActive("cancelled") + return nil + case <-ticker.C: + } + continue + } + + r.runBatch(ctx, due) + } +} + +// runBatch executes the due tasks concurrently up to Concurrency. Tasks that +// become non-pending while queued are skipped by executeOne's re-check. +func (r *TaskRunner) runBatch(ctx context.Context, tasks []*Task) { + sem := make(chan struct{}, r.opts.Concurrency) + var wg sync.WaitGroup + for _, task := range tasks { + if r.terminalCount() >= r.opts.MaxTotalTasks { + break + } + select { + case <-ctx.Done(): + return + case sem <- struct{}{}: + } + wg.Add(1) + go func(t *Task) { + defer wg.Done() + defer func() { <-sem }() + r.executeOne(ctx, t) + }(task) + } + wg.Wait() +} + +func (r *TaskRunner) executeOne(ctx context.Context, task *Task) { + // Watchdog: never start work once the terminal-state cap is reached. This + // guard lives here (not just at batch build time) because goroutines are + // queued before earlier ones complete. + if r.terminalCount() >= r.opts.MaxTotalTasks { + return + } + cur, ok := r.store.Get(task.ID) + if !ok || cur.Status != TaskStatusPending { + return + } + + r.store.Update(task.ID, func(t *Task) { + t.Status = TaskStatusInProgress + if t.Owner == "" { + t.Owner = "task-runner" + } + if t.Metadata == nil { + t.Metadata = map[string]any{} + } + t.Metadata["execStartedAt"] = time.Now().UTC().Format(time.RFC3339Nano) + }) + r.progress(task, EventStarted) + r.bump("executed") + + execCtx := ctx + cancel := func() {} + if r.opts.DefaultTimeout > 0 { + execCtx, cancel = context.WithTimeout(ctx, r.opts.DefaultTimeout) + } + out, err := r.opts.Execute(execCtx, task) + cancel() + + if err == nil { + r.store.Update(task.ID, func(t *Task) { + t.Status = TaskStatusCompleted + if t.Checkpoint == nil { + t.Checkpoint = map[string]any{} + } + t.Checkpoint["result"] = out + if t.Metadata == nil { + t.Metadata = map[string]any{} + } + t.Metadata["execFinishedAt"] = time.Now().UTC().Format(time.RFC3339Nano) + }) + r.progress(task, EventCompleted) + r.bump("completed") + return + } + + // The run itself was cancelled/stopped (not a per-task timeout): park the + // task as cancelled rather than consuming retry budget for an aborted run. + if ctx.Err() != nil { + if ok, _ := r.store.Cancel(task.ID, "run stopped"); ok { + r.bump("cancelled") + } + r.progress(task, EventCancelled) + return + } + + requeued, err := r.store.MarkFailed(task.ID, err.Error()) + if err != nil { + r.progress(task, EventFailed) + r.bump("failed") + return + } + if requeued { + r.progress(task, EventRetrying) + return + } + if r.tryReplan(ctx, task) { + r.progress(task, EventReplanned) + return + } + r.progress(task, EventFailed) + r.bump("failed") +} + +// tryReplan requeues a failed task with a fresh budget when the host's +// OnReplan hook approves and the per-task replan cap is not exhausted. +func (r *TaskRunner) tryReplan(ctx context.Context, task *Task) bool { + if r.opts.OnReplan == nil { + return false + } + r.mu.Lock() + n := r.replansByTask[task.ID] + r.mu.Unlock() + if n >= r.opts.MaxReplans { + return false + } + ok, err := r.opts.OnReplan(ctx, task) + if err != nil || !ok { + return false + } + if _, err := r.store.Requeue(task.ID); err != nil { + return false + } + r.mu.Lock() + r.replansByTask[task.ID]++ + r.replanned++ + r.mu.Unlock() + return true +} + +// dueTasks filters ready tasks to those whose retry backoff has elapsed. A task +// with no pending backoff tick is immediately due. +func (r *TaskRunner) dueTasks(tasks []*Task, now time.Time) []*Task { + var due []*Task + for _, t := range tasks { + tick := metaInt(t.Metadata, "retryBackoffTick") + if tick <= 0 { + due = append(due, t) + continue + } + if !now.Before(t.UpdatedAt.Add(r.backoff(tick))) { + due = append(due, t) + } + } + return due +} + +// quiescent reports whether no ready work exists and nothing is in flight. +// Tasks blocked behind failed/skipped/cancelled dependencies are not "active": +// they never become ready, so they must not keep the run alive forever. +func (r *TaskRunner) quiescent() bool { + if len(r.store.GetReadyWork()) > 0 { + return false + } + for _, t := range r.store.List() { + switch t.Status { + case TaskStatusInProgress, TaskStatusReviewing: + return false + } + } + return true +} + +// cancelActive parks all in-flight tasks as cancelled with the given reason. +func (r *TaskRunner) cancelActive(reason string) { + for _, t := range r.store.List() { + if t.Status != TaskStatusInProgress && t.Status != TaskStatusReviewing { + continue + } + if ok, _ := r.store.Cancel(t.ID, reason); ok { + r.bump("cancelled") + } + } +} + +func (r *TaskRunner) terminalCount() int { + r.mu.Lock() + defer r.mu.Unlock() + return r.completed + r.failed + r.skipped + r.cancelled +} + +func (r *TaskRunner) bump(kind string) { + r.mu.Lock() + defer r.mu.Unlock() + switch kind { + case "executed": + r.executed++ + case "completed": + r.completed++ + case "failed": + r.failed++ + case "skipped": + r.skipped++ + case "cancelled": + r.cancelled++ + case "replanned": + r.replanned++ + } +} + +func (r *TaskRunner) progress(task *Task, event TaskRunnerEvent) { + if r.opts.OnProgress != nil { + r.opts.OnProgress(task, event) + } +} + +// Start launches Run in a background goroutine. Stop cancels it; Wait blocks +// until it finishes. +func (r *TaskRunner) Start(ctx context.Context) { + r.mu.Lock() + if r.started { + r.mu.Unlock() + return + } + r.started = true + r.finished = false + runCtx, cancel := context.WithCancel(ctx) + r.cancel = cancel + r.doneCh = make(chan struct{}) + r.mu.Unlock() + + go func() { + defer func() { + r.mu.Lock() + r.finished = true + r.mu.Unlock() + close(r.doneCh) + }() + _ = r.Run(runCtx) + }() +} + +// Stop cancels a background run. Idempotent; safe when never started. +func (r *TaskRunner) Stop() { + r.mu.Lock() + c := r.cancel + r.mu.Unlock() + if c != nil { + c() + } +} + +// Wait blocks until the background run finishes or ctx is done. +func (r *TaskRunner) Wait(ctx context.Context) error { + r.mu.Lock() + done := r.doneCh + r.mu.Unlock() + if done == nil { + return fmt.Errorf("task runner: not started") + } + select { + case <-done: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +// Status returns a point-in-time snapshot of runner activity. +func (r *TaskRunner) Status() TaskRunnerStats { + r.mu.Lock() + defer r.mu.Unlock() + return TaskRunnerStats{ + Running: r.started && !r.finished, + Executed: r.executed, + Completed: r.completed, + Failed: r.failed, + Skipped: r.skipped, + Cancelled: r.cancelled, + Replanned: r.replanned, + } +} + +// metaInt reads an integer metadata value that may have been round-tripped +// through JSON (float64) or set directly in memory (int). +func metaInt(m map[string]any, key string) int { + if m == nil { + return 0 + } + switch v := m[key].(type) { + case int: + return v + case int64: + return int(v) + case float64: + return int(v) + default: + return 0 + } +} diff --git a/internal/tool/task_executor_test.go b/internal/tool/task_executor_test.go new file mode 100644 index 00000000..0539e0f9 --- /dev/null +++ b/internal/tool/task_executor_test.go @@ -0,0 +1,402 @@ +package tool + +import ( + "context" + "fmt" + "strings" + "sync" + "testing" + "time" +) + +// fakeTaskExecutor records the tasks it runs and fails those on the fail list. +type fakeTaskExecutor struct { + mu sync.Mutex + ran []string + fail map[string]error + delay time.Duration + output map[string]string +} + +func newFakeExecutor(fail map[string]error) *fakeTaskExecutor { + return &fakeTaskExecutor{ + fail: fail, + output: map[string]string{}, + } +} + +func (f *fakeTaskExecutor) execute(ctx context.Context, t *Task) (string, error) { + f.mu.Lock() + f.ran = append(f.ran, t.ID) + err := f.fail[t.ID] + out := f.output[t.ID] + f.mu.Unlock() + if ctx != nil { + select { + case <-ctx.Done(): + return "", ctx.Err() + default: + } + } + if f.delay > 0 { + select { + case <-time.After(f.delay): + case <-ctx.Done(): + return "", ctx.Err() + } + } + if err != nil { + return "", err + } + if out == "" { + out = "done:" + t.ID + } + return out, nil +} + +func newStoreWith(tasks ...func(*Task)) *TaskStore { + s := &TaskStore{tasks: make(map[string]*Task)} + for _, setup := range tasks { + t := s.Create("subject", "description", "", nil) + setup(t) + } + return s +} + +func TestTaskRunnerCompletesReadyWork(t *testing.T) { + s := &TaskStore{tasks: make(map[string]*Task)} + a := s.Create("a", "do a", "", nil) + b := s.Create("b", "do b", "", nil) + + exec := newFakeExecutor(nil) + r := NewTaskRunner(TaskRunnerOptions{ + Store: s, + Execute: exec.execute, + PollInterval: time.Millisecond, + }) + if err := r.Run(context.Background()); err != nil { + t.Fatal(err) + } + + got, _ := s.Get(a.ID) + if got.Status != TaskStatusCompleted { + t.Fatalf("task a status = %q, want completed", got.Status) + } + gotB, _ := s.Get(b.ID) + if gotB.Status != TaskStatusCompleted { + t.Fatalf("task b status = %q, want completed", gotB.Status) + } + st := r.Status() + if st.Completed != 2 || st.Executed != 2 { + t.Fatalf("unexpected stats: %+v", st) + } + if got.Checkpoint["result"] != "done:task_1" { + t.Fatalf("expected result checkpoint, got %+v", got.Checkpoint) + } +} + +func TestTaskRunnerRespectsDependencies(t *testing.T) { + s := &TaskStore{tasks: make(map[string]*Task)} + a := s.Create("a", "do a", "", nil) + b := s.Create("b", "do b", "", nil) + s.Update(b.ID, func(t *Task) { + t.Dependencies = []TaskDependency{{TargetID: a.ID, Type: "blocks"}} + }) + + exec := newFakeExecutor(nil) + r := NewTaskRunner(TaskRunnerOptions{ + Store: s, + Execute: exec.execute, + PollInterval: time.Millisecond, + }) + if err := r.Run(context.Background()); err != nil { + t.Fatal(err) + } + + exec.mu.Lock() + defer exec.mu.Unlock() + if len(exec.ran) != 2 { + t.Fatalf("expected 2 executions, got %v", exec.ran) + } + if exec.ran[0] != a.ID || exec.ran[1] != b.ID { + t.Fatalf("dependency order violated: %v", exec.ran) + } +} + +func TestTaskRunnerRetriesWithinBudget(t *testing.T) { + s := &TaskStore{tasks: make(map[string]*Task)} + task := s.Create("flaky", "do flaky", "", nil) + s.Update(task.ID, func(t *Task) { t.MaxAttempts = 3 }) + + // Fail the first two attempts, succeed on the third. + attempts := map[string]int{} + var mu sync.Mutex + exec := func(_ context.Context, t *Task) (string, error) { + mu.Lock() + attempts[t.ID]++ + n := attempts[t.ID] + mu.Unlock() + if n < 3 { + return "", fmt.Errorf("transient failure %d", n) + } + return "ok", nil + } + + r := NewTaskRunner(TaskRunnerOptions{ + Store: s, + Execute: exec, + Backoff: func(int) time.Duration { return time.Millisecond }, + PollInterval: time.Millisecond, + }) + if err := r.Run(context.Background()); err != nil { + t.Fatal(err) + } + + got, _ := s.Get(task.ID) + if got.Status != TaskStatusCompleted { + t.Fatalf("expected completed after retries, got %q (lastError=%q)", got.Status, got.LastError) + } + if got.Attempts != 2 { + t.Fatalf("expected 2 failed attempts recorded, got %d", got.Attempts) + } +} + +func TestTaskRunnerParksFailedAfterBudget(t *testing.T) { + s := &TaskStore{tasks: make(map[string]*Task)} + task := s.Create("doomed", "do doomed", "", nil) + s.Update(task.ID, func(t *Task) { t.MaxAttempts = 2 }) + + exec := func(_ context.Context, t *Task) (string, error) { + return "", fmt.Errorf("always fails") + } + + r := NewTaskRunner(TaskRunnerOptions{ + Store: s, + Execute: exec, + Backoff: func(int) time.Duration { return time.Millisecond }, + PollInterval: time.Millisecond, + }) + if err := r.Run(context.Background()); err != nil { + t.Fatal(err) + } + + got, _ := s.Get(task.ID) + if got.Status != TaskStatusFailed { + t.Fatalf("expected failed after budget, got %q", got.Status) + } + if got.Attempts != 2 { + t.Fatalf("expected 2 attempts, got %d", got.Attempts) + } + st := r.Status() + if st.Failed != 1 { + t.Fatalf("expected 1 failed in stats, got %+v", st) + } +} + +func TestTaskRunnerReplansFailedTask(t *testing.T) { + s := &TaskStore{tasks: make(map[string]*Task)} + task := s.Create("replan", "do replan", "", nil) + s.Update(task.ID, func(t *Task) { t.MaxAttempts = 1 }) + + // Fail the first round (budget 1), then replan requeues it; succeed after. + round := 0 + exec := func(_ context.Context, t *Task) (string, error) { + if round == 0 { + round = 1 + return "", fmt.Errorf("first round fails") + } + return "ok", nil + } + replans := 0 + r := NewTaskRunner(TaskRunnerOptions{ + Store: s, + Execute: exec, + Backoff: func(int) time.Duration { return time.Millisecond }, + PollInterval: time.Millisecond, + OnReplan: func(_ context.Context, t *Task) (bool, error) { + replans++ + return true, nil + }, + }) + if err := r.Run(context.Background()); err != nil { + t.Fatal(err) + } + + got, _ := s.Get(task.ID) + if got.Status != TaskStatusCompleted { + t.Fatalf("expected completed after replan, got %q", got.Status) + } + if replans != 1 { + t.Fatalf("expected 1 replan, got %d", replans) + } + st := r.Status() + if st.Replanned != 1 { + t.Fatalf("expected 1 replanned in stats, got %+v", st) + } +} + +func TestTaskRunnerWatchdogStops(t *testing.T) { + s := &TaskStore{tasks: make(map[string]*Task)} + for i := 0; i < 10; i++ { + s.Create(fmt.Sprintf("t%d", i), "desc", "", nil) + } + + exec := newFakeExecutor(nil) + r := NewTaskRunner(TaskRunnerOptions{ + Store: s, + Execute: exec.execute, + PollInterval: time.Millisecond, + MaxTotalTasks: 3, + }) + if err := r.Run(context.Background()); err != nil { + t.Fatal(err) + } + + exec.mu.Lock() + n := len(exec.ran) + exec.mu.Unlock() + if n != 3 { + t.Fatalf("watchdog should stop after 3 tasks, ran %d", n) + } + st := r.Status() + if st.Completed != 3 { + t.Fatalf("expected 3 completed, got %+v", st) + } + // Remaining tasks stay pending. + remaining := 0 + for _, t := range s.List() { + if t.Status == TaskStatusPending { + remaining++ + } + } + if remaining != 7 { + t.Fatalf("expected 7 pending remaining, got %d", remaining) + } +} + +func TestTaskRunnerCancelsInFlightOnStop(t *testing.T) { + s := &TaskStore{tasks: make(map[string]*Task)} + task := s.Create("slow", "slow task", "", nil) + + started := make(chan struct{}) + exec := func(ctx context.Context, t *Task) (string, error) { + close(started) + <-ctx.Done() + return "", ctx.Err() + } + + r := NewTaskRunner(TaskRunnerOptions{ + Store: s, + Execute: exec, + PollInterval: time.Millisecond, + DefaultTimeout: 5 * time.Second, + }) + ctx := context.Background() + r.Start(ctx) + <-started + r.Stop() + if err := r.Wait(context.Background()); err != nil { + t.Fatal(err) + } + + got, _ := s.Get(task.ID) + if got.Status != TaskStatusCancelled { + t.Fatalf("expected cancelled on stop, got %q", got.Status) + } +} + +func TestTaskRunnerQuiescesBehindFailedBlocker(t *testing.T) { + s := &TaskStore{tasks: make(map[string]*Task)} + a := s.Create("a", "do a", "", nil) + b := s.Create("b", "depends on a", "", nil) + s.Update(b.ID, func(t *Task) { t.Dependencies = []TaskDependency{{TargetID: a.ID, Type: "blocks"}} }) + s.Update(a.ID, func(t *Task) { t.MaxAttempts = 1 }) + // Mark a failed so b can never become ready. + s.MarkFailed(a.ID, "cannot do a") + + exec := newFakeExecutor(nil) + r := NewTaskRunner(TaskRunnerOptions{ + Store: s, + Execute: exec.execute, + PollInterval: time.Millisecond, + }) + // Must terminate rather than spin forever behind the failed blocker. + done := make(chan error, 1) + go func() { done <- r.Run(context.Background()) }() + select { + case err := <-done: + if err != nil { + t.Fatal(err) + } + case <-time.After(3 * time.Second): + t.Fatal("runner spun forever behind a failed blocker") + } +} + +func TestTaskRunnerProgressEvents(t *testing.T) { + s := &TaskStore{tasks: make(map[string]*Task)} + s.Create("a", "do a", "", nil) + + var events []string + var mu sync.Mutex + exec := newFakeExecutor(nil) + r := NewTaskRunner(TaskRunnerOptions{ + Store: s, + Execute: exec.execute, + PollInterval: time.Millisecond, + OnProgress: func(t *Task, e TaskRunnerEvent) { + mu.Lock() + events = append(events, string(e)) + mu.Unlock() + }, + }) + if err := r.Run(context.Background()); err != nil { + t.Fatal(err) + } + + mu.Lock() + defer mu.Unlock() + joined := strings.Join(events, ",") + if !strings.Contains(joined, string(EventStarted)) || !strings.Contains(joined, string(EventCompleted)) { + t.Fatalf("expected started+completed events, got %q", joined) + } +} + +func TestSkipAndCancel(t *testing.T) { + s := &TaskStore{tasks: make(map[string]*Task)} + task := s.Create("a", "do a", "", nil) + + if ok, _ := s.Skip(task.ID, "not needed"); !ok { + t.Fatal("skip should succeed") + } + got, _ := s.Get(task.ID) + if got.Status != TaskStatusSkipped || got.Metadata["skipReason"] != "not needed" { + t.Fatalf("unexpected after skip: %+v", got) + } + // Skipping a completed task is a no-op. + s.Update(task.ID, func(t *Task) { t.Status = TaskStatusPending }) + if ok, _ := s.Skip(task.ID, "x"); !ok { + t.Fatal("re-skip of pending should succeed") + } + s.Update(task.ID, func(t *Task) { t.Status = TaskStatusCompleted }) + if ok, _ := s.Cancel(task.ID, "too late"); ok { + t.Fatal("cancelling a completed task should be a no-op") + } + + c := s.Create("c", "do c", "", nil) + if ok, _ := s.Cancel(c.ID, "aborted"); !ok { + t.Fatal("cancel should succeed") + } + gotC, _ := s.Get(c.ID) + if gotC.Status != TaskStatusCancelled || gotC.Metadata["cancelReason"] != "aborted" { + t.Fatalf("unexpected after cancel: %+v", gotC) + } +} + +func TestTaskRunnerMissingExecute(t *testing.T) { + s := &TaskStore{tasks: make(map[string]*Task)} + r := NewTaskRunner(TaskRunnerOptions{Store: s}) + if err := r.Run(context.Background()); err == nil { + t.Fatal("expected error when execute is missing") + } +} diff --git a/internal/tool/task_run_tool.go b/internal/tool/task_run_tool.go new file mode 100644 index 00000000..ed5ccf58 --- /dev/null +++ b/internal/tool/task_run_tool.go @@ -0,0 +1,83 @@ +package tool + +import ( + "context" + "encoding/json" + "fmt" + "strings" + "time" +) + +// TaskRunTool drives ready tasks through the task executor. It lets the agent +// (or a host) turn the validated task graph into execution: tasks run in +// dependency order, are retried up to their retry budget with backoff, and are +// parked failed when the budget is exhausted. It is the tool-level front door +// for TaskRunner. +type TaskRunTool struct{} + +func (TaskRunTool) Name() string { return "TaskRun" } +func (TaskRunTool) Aliases() []string { return []string{"task_run"} } +func (TaskRunTool) Description() string { + return "Execute all ready tasks (pending with no blockers) through the task executor. " + + "Tasks run in dependency order, are retried up to their retry budget, and are parked failed " + + "when the budget is exhausted. Returns a run summary." +} + +func (TaskRunTool) Parameters() map[string]interface{} { + return map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "timeout_sec": map[string]interface{}{ + "type": "integer", + "description": "Per-task execution timeout in seconds (default 300)", + }, + "max_total_tasks": map[string]interface{}{ + "type": "integer", + "description": "Watchdog cap on distinct tasks that may reach a terminal state (default 50)", + }, + }, + } +} + +func (TaskRunTool) Execute(ctx context.Context, input json.RawMessage) (string, error) { + tc := GetToolContext(ctx) + if tc == nil || tc.TaskExecutor == nil { + return "", fmt.Errorf("TaskRun requires a task executor; none is configured for this session") + } + + var p struct { + TimeoutSec int `json:"timeout_sec"` + MaxTotalTasks int `json:"max_total_tasks"` + } + if input != nil { + _ = json.Unmarshal(input, &p) + } + timeout := time.Duration(p.TimeoutSec) * time.Second + if timeout <= 0 { + timeout = 300 * time.Second + } + + runner := NewTaskRunner(TaskRunnerOptions{ + Store: GetTaskStore(), + Execute: tc.TaskExecutor, + DefaultTimeout: timeout, + MaxTotalTasks: p.MaxTotalTasks, + }) + if err := runner.Run(ctx); err != nil { + return "", err + } + s := runner.Status() + parts := []string{ + fmt.Sprintf("%d executed", s.Executed), + fmt.Sprintf("%d completed", s.Completed), + fmt.Sprintf("%d failed", s.Failed), + fmt.Sprintf("%d replanned", s.Replanned), + } + if s.Skipped > 0 { + parts = append(parts, fmt.Sprintf("%d skipped", s.Skipped)) + } + if s.Cancelled > 0 { + parts = append(parts, fmt.Sprintf("%d cancelled", s.Cancelled)) + } + return "Task run finished: " + strings.Join(parts, ", "), nil +} diff --git a/internal/tool/task_run_tool_test.go b/internal/tool/task_run_tool_test.go new file mode 100644 index 00000000..ff4a68ac --- /dev/null +++ b/internal/tool/task_run_tool_test.go @@ -0,0 +1,69 @@ +package tool + +import ( + "context" + "strings" + "testing" + "time" +) + +func TestTaskRunToolExecutesReadyTasks(t *testing.T) { + s := &TaskStore{tasks: make(map[string]*Task)} + s.Create("a", "do a", "", nil) + s.Create("b", "do b", "", nil) + // Replace the global store reference used by the tool for the test. + prev := globalTaskStore + globalTaskStore = s + defer func() { globalTaskStore = prev }() + + exec := newFakeExecutor(nil) + ctx := WithToolContext(context.Background(), &ToolContext{TaskExecutor: exec.execute}) + out, err := TaskRunTool{}.Execute(ctx, nil) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if !strings.Contains(out, "2 completed") { + t.Fatalf("expected 2 completed in summary, got: %s", out) + } + got, _ := s.Get("task_1") + if got.Status != TaskStatusCompleted { + t.Fatalf("task_1 status = %q, want completed", got.Status) + } +} + +func TestTaskRunToolRequiresExecutor(t *testing.T) { + ctx := WithToolContext(context.Background(), &ToolContext{}) + _, err := TaskRunTool{}.Execute(ctx, nil) + if err == nil { + t.Fatal("expected error when no executor configured") + } + if !strings.Contains(err.Error(), "task executor") { + t.Fatalf("unexpected error: %v", err) + } +} + +func TestTaskRunToolCancelsOnContextTimeout(t *testing.T) { + s := &TaskStore{tasks: make(map[string]*Task)} + s.Create("slow", "slow task", "", nil) + prev := globalTaskStore + globalTaskStore = s + defer func() { globalTaskStore = prev }() + + exec := func(ctx context.Context, t *Task) (string, error) { + <-ctx.Done() + return "", ctx.Err() + } + + runCtx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + runCtx = WithToolContext(runCtx, &ToolContext{TaskExecutor: exec}) + _, err := TaskRunTool{}.Execute(runCtx, nil) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + got, _ := s.Get("task_1") + if got.Status != TaskStatusCancelled { + t.Fatalf("expected task cancelled on ctx timeout, got %q", got.Status) + } +} diff --git a/internal/tool/task_runner.go b/internal/tool/task_runner.go index 8cd5800f..59c6adae 100644 --- a/internal/tool/task_runner.go +++ b/internal/tool/task_runner.go @@ -213,6 +213,64 @@ func (s *TaskStore) Requeue(id string) (bool, error) { return true, nil } +// Skip parks a task as skipped without executing it. The reason is recorded in +// metadata for later inspection (skipReason). +func (s *TaskStore) Skip(id, reason string) (bool, error) { + s.mu.Lock() + t, ok := s.tasks[id] + if !ok { + s.mu.Unlock() + return false, fmt.Errorf("task %q not found", id) + } + if t.Status == TaskStatusCompleted || t.Status == TaskStatusSkipped { + s.mu.Unlock() + return false, nil + } + t.Status = TaskStatusSkipped + if t.Metadata == nil { + t.Metadata = make(map[string]any) + } + if reason != "" { + t.Metadata["skipReason"] = reason + } + t.UpdatedAt = time.Now() + persist := s.persist + s.mu.Unlock() + if persist != nil { + _ = s.Save("") + } + return true, nil +} + +// Cancel parks a task as cancelled (e.g. the run was aborted). The reason is +// recorded in metadata for later inspection (cancelReason). +func (s *TaskStore) Cancel(id, reason string) (bool, error) { + s.mu.Lock() + t, ok := s.tasks[id] + if !ok { + s.mu.Unlock() + return false, fmt.Errorf("task %q not found", id) + } + if t.Status == TaskStatusCompleted || t.Status == TaskStatusCancelled { + s.mu.Unlock() + return false, nil + } + t.Status = TaskStatusCancelled + if t.Metadata == nil { + t.Metadata = make(map[string]any) + } + if reason != "" { + t.Metadata["cancelReason"] = reason + } + t.UpdatedAt = time.Now() + persist := s.persist + s.mu.Unlock() + if persist != nil { + _ = s.Save("") + } + return true, nil +} + // Checkpoint merges resumable progress onto a task without changing its // status. A replan or resume reads the checkpoint to avoid starting from zero. func (s *TaskStore) Checkpoint(id string, data map[string]any) (bool, error) { diff --git a/internal/tool/task_schedule.go b/internal/tool/task_schedule.go index 826936cd..055141b5 100644 --- a/internal/tool/task_schedule.go +++ b/internal/tool/task_schedule.go @@ -41,7 +41,8 @@ func (s *TaskStore) Schedule() (TaskSchedule, error) { return TaskSchedule{}, fmt.Errorf("task schedule contains invalid task identity %q", id) } switch task.Status { - case TaskStatusPending, TaskStatusInProgress, TaskStatusCompleted, TaskStatusFailed: + case TaskStatusPending, TaskStatusInProgress, TaskStatusReviewing, + TaskStatusCompleted, TaskStatusFailed, TaskStatusSkipped, TaskStatusCancelled: default: return TaskSchedule{}, fmt.Errorf("task %q has invalid status %q", id, task.Status) } diff --git a/internal/tool/tool.go b/internal/tool/tool.go index 54696db9..ceff3e81 100644 --- a/internal/tool/tool.go +++ b/internal/tool/tool.go @@ -97,6 +97,9 @@ type ToolContext struct { // WorkingDir, when set, is used as cmd.Dir for Bash and as the preferred // workspace root for path tools (subagent worktree isolation). WorkingDir string + // TaskExecutor, when set, arms the TaskRun tool: it runs one task from the + // store (e.g. by spawning a sub-agent). Nil disables TaskRun. + TaskExecutor TaskExecutorFunc // Lint configures the optional post-write auto-lint cycle. The zero value // (Enabled=false) keeps linting off so users are not surprised. Lint lint.Config