From 627b0b1c32e2412bddc1a7c6531502f2b6263c37 Mon Sep 17 00:00:00 2001 From: dj-oyu <68707227+dj-oyu@users.noreply.github.com> Date: Sun, 22 Feb 2026 07:03:48 +0900 Subject: [PATCH] feat: auto-continue plan execution after step completion The plan nudge was background-only and only fired when no progress was made. Now it fires for all plan executions (foreground included) and whenever unchecked steps remain, with a context-aware message depending on whether progress was recorded. Co-Authored-By: Claude Opus 4.6 --- pkg/agent/loop.go | 33 +++--- pkg/agent/loop_test.go | 237 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 256 insertions(+), 14 deletions(-) diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index db12a97f5..2d38750f7 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -1352,7 +1352,7 @@ func (al *AgentLoop) runLLMIteration( // Snapshot unchecked step count before tool loop so we can detect progress. preUnchecked := -1 // -1 = not tracking - if opts.Background && agent.ContextBuilder.GetPlanStatus() == "executing" { + if agent.ContextBuilder.GetPlanStatus() == "executing" { preUnchecked = strings.Count(agent.ContextBuilder.ReadMemory(), "- [ ]") } @@ -1506,31 +1506,36 @@ func (al *AgentLoop) runLLMIteration( // Check if no tool calls - we're done if len(response.ToolCalls) == 0 { - // Background plan continuation: if unchecked steps remain and - // none were marked during this heartbeat, nudge the LLM to + // Plan continuation: if unchecked steps remain, nudge the LLM to // either mark completed steps or continue working on them. - // This serves as both a marking reminder and a continuation trigger - // (otherwise remaining steps wait until the next heartbeat). + // This fires for both foreground and background plan execution, + // ensuring the loop doesn't exit prematurely after marking a step. curUnchecked := 0 if preUnchecked > 0 { curUnchecked = strings.Count(agent.ContextBuilder.ReadMemory(), "- [ ]") } - if preUnchecked > 0 && !planMarkNudged && - agent.ContextBuilder.GetPlanStatus() == "executing" && - curUnchecked == preUnchecked { + if curUnchecked > 0 && !planMarkNudged && + agent.ContextBuilder.GetPlanStatus() == "executing" { planMarkNudged = true messages = append(messages, providers.Message{ Role: "assistant", Content: response.Content, }) - messages = append(messages, providers.Message{ - Role: "user", - Content: fmt.Sprintf("[System] %d unchecked steps remain in MEMORY.md and none were marked [x] during this session. "+ + var nudgeMsg string + if curUnchecked == preUnchecked { + nudgeMsg = fmt.Sprintf("[System] %d unchecked steps remain in MEMORY.md and "+ + "none were marked [x] during this session. "+ "If you completed any steps, use edit_file to mark them [x] now. "+ - "If steps are still in progress, continue working on them.", - curUnchecked), + "If steps are still in progress, continue working on them.", curUnchecked) + } else { + nudgeMsg = fmt.Sprintf("[System] Progress recorded. %d unchecked steps remain. "+ + "Continue working on the next step.", curUnchecked) + } + messages = append(messages, providers.Message{ + Role: "user", + Content: nudgeMsg, }) - logger.InfoCF("agent", "Nudging background task: mark or continue plan steps", + logger.InfoCF("agent", "Nudging plan execution: continue plan steps", map[string]any{"agent_id": agent.ID, "iteration": iteration, "unchecked": curUnchecked}) continue } diff --git a/pkg/agent/loop_test.go b/pkg/agent/loop_test.go index e27fa7d5f..72e7d7ec4 100644 --- a/pkg/agent/loop_test.go +++ b/pkg/agent/loop_test.go @@ -1852,3 +1852,240 @@ func TestSanitizeHistoryForProvider_MultiToolCall(t *testing.T) { t.Errorf("expected 2 tool results, got %d", toolCount) } } + +// ---------- plan nudge tests ---------- + +// countingMockProvider counts Chat calls and always returns text-only responses. +type countingMockProvider struct { + callCount int +} + +func (m *countingMockProvider) Chat( + ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + opts map[string]any, +) (*providers.LLMResponse, error) { + m.callCount++ + return &providers.LLMResponse{ + Content: fmt.Sprintf("Response %d", m.callCount), + ToolCalls: []providers.ToolCall{}, + }, nil +} + +func (m *countingMockProvider) GetDefaultModel() string { + return "mock-counting-model" +} + +func TestPlanNudge_ForegroundExecution(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "agent-nudge-test-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + Model: "test-model", + MaxTokens: 4096, + MaxToolIterations: 10, + }, + }, + } + + provider := &countingMockProvider{} + msgBus := bus.NewMessageBus() + al := NewAgentLoop(cfg, msgBus, provider) + + agent := al.registry.GetDefaultAgent() + if agent == nil { + t.Fatal("no default agent") + } + + // Write a plan in executing status with unchecked steps + plan := "# Active Plan\n\n> Task: Test\n> Status: executing\n> Phase: 1\n\n## Phase 1: Setup\n- [ ] Step one\n- [ ] Step two\n\n## Context\n" + agent.ContextBuilder.WriteMemory(plan) + + // Process a foreground message (no background metadata) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + msg := bus.InboundMessage{ + Channel: "test", + SenderID: "user1", + ChatID: "chat1", + Content: "continue working", + SessionKey: "nudge-test", + } + _, err = al.processMessage(ctx, msg) + if err != nil { + t.Fatalf("processMessage failed: %v", err) + } + + // The provider should have been called at least 2 times: + // 1st call: returns text-only → nudge fires (unchecked steps remain) + // 2nd call: returns text-only → nudge already fired, loop exits + if provider.callCount < 2 { + t.Errorf("expected at least 2 provider calls (nudge should trigger continuation), got %d", provider.callCount) + } +} + +func TestPlanNudge_NoNudgeWhenAllStepsComplete(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "agent-nudge-test-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + Model: "test-model", + MaxTokens: 4096, + MaxToolIterations: 10, + }, + }, + } + + provider := &countingMockProvider{} + msgBus := bus.NewMessageBus() + al := NewAgentLoop(cfg, msgBus, provider) + + agent := al.registry.GetDefaultAgent() + if agent == nil { + t.Fatal("no default agent") + } + + // Write a plan where all steps are already checked + plan := "# Active Plan\n\n> Task: Test\n> Status: executing\n> Phase: 1\n\n## Phase 1: Setup\n- [x] Step one\n- [x] Step two\n\n## Context\n" + agent.ContextBuilder.WriteMemory(plan) + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + msg := bus.InboundMessage{ + Channel: "test", + SenderID: "user1", + ChatID: "chat1", + Content: "all done", + SessionKey: "nudge-test-complete", + } + _, err = al.processMessage(ctx, msg) + if err != nil { + t.Fatalf("processMessage failed: %v", err) + } + + // No unchecked steps → preUnchecked=0 → no nudge → only 1 provider call + if provider.callCount != 1 { + t.Errorf("expected exactly 1 provider call (no nudge needed), got %d", provider.callCount) + } +} + +func TestPlanNudge_ProgressMessage(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "agent-nudge-test-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + Model: "test-model", + MaxTokens: 4096, + MaxToolIterations: 10, + }, + }, + } + + // Provider that checks the nudge message content on the 2nd call + var nudgeContent string + provider := &nudgeCaptureMockProvider{onSecondCall: func(msgs []providers.Message) { + // The last user message should be the nudge + for i := len(msgs) - 1; i >= 0; i-- { + if msgs[i].Role == "user" { + nudgeContent = msgs[i].Content + break + } + } + }} + msgBus := bus.NewMessageBus() + al := NewAgentLoop(cfg, msgBus, provider) + + agent := al.registry.GetDefaultAgent() + if agent == nil { + t.Fatal("no default agent") + } + + // Write a plan with 3 unchecked steps; the provider edits memory to + // mark one step between calls (simulated by the first-call hook). + plan := "# Active Plan\n\n> Task: Test\n> Status: executing\n> Phase: 1\n\n## Phase 1: Setup\n- [ ] Step one\n- [ ] Step two\n- [ ] Step three\n\n## Context\n" + agent.ContextBuilder.WriteMemory(plan) + + // After the first LLM response (no tool calls), simulate that + // one step was marked [x] externally (as if the AI did it via tool). + // We do this by hooking the provider's first call to mutate memory. + provider.onFirstCall = func() { + updated := strings.Replace(agent.ContextBuilder.ReadMemory(), "- [ ] Step one", "- [x] Step one", 1) + agent.ContextBuilder.WriteMemory(updated) + } + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + msg := bus.InboundMessage{ + Channel: "test", + SenderID: "user1", + ChatID: "chat1", + Content: "work on the plan", + SessionKey: "nudge-progress-test", + } + _, err = al.processMessage(ctx, msg) + if err != nil { + t.Fatalf("processMessage failed: %v", err) + } + + // Should have gotten the "Progress recorded" nudge (not the "none were marked" one) + if !strings.Contains(nudgeContent, "Progress recorded") { + t.Errorf("expected 'Progress recorded' nudge, got %q", nudgeContent) + } + if !strings.Contains(nudgeContent, "2 unchecked steps remain") { + t.Errorf("expected '2 unchecked steps remain' in nudge, got %q", nudgeContent) + } +} + +// nudgeCaptureMockProvider calls hooks on 1st and 2nd Chat invocations. +type nudgeCaptureMockProvider struct { + callCount int + onFirstCall func() + onSecondCall func([]providers.Message) +} + +func (m *nudgeCaptureMockProvider) Chat( + ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + opts map[string]any, +) (*providers.LLMResponse, error) { + m.callCount++ + if m.callCount == 1 && m.onFirstCall != nil { + m.onFirstCall() + } + if m.callCount == 2 && m.onSecondCall != nil { + m.onSecondCall(messages) + } + return &providers.LLMResponse{ + Content: fmt.Sprintf("Response %d", m.callCount), + ToolCalls: []providers.ToolCall{}, + }, nil +} + +func (m *nudgeCaptureMockProvider) GetDefaultModel() string { + return "mock-nudge-model" +}