From 5fe734e2aa29de067f6b56176e957dfb30ccd9ff Mon Sep 17 00:00:00 2001 From: dj-oyu <68707227+dj-oyu@users.noreply.github.com> Date: Fri, 20 Feb 2026 14:07:41 +0900 Subject: [PATCH] feat(agent): add in-loop task reminder to prevent focus drift Inject ephemeral user-role reminders every N iterations in the tool call loop so the LLM retains the original task instruction even when tool results flood the message array. Includes blocker tracking from failed tool results. Co-Authored-By: Claude Opus 4.6 --- pkg/agent/instance.go | 15 +++++-- pkg/agent/loop.go | 52 ++++++++++++++++++++++++ pkg/agent/loop_test.go | 91 ++++++++++++++++++++++++++++++++++++++++++ pkg/config/config.go | 6 ++- 4 files changed, 158 insertions(+), 6 deletions(-) diff --git a/pkg/agent/instance.go b/pkg/agent/instance.go index 54a5396e7..83a3cc0a6 100644 --- a/pkg/agent/instance.go +++ b/pkg/agent/instance.go @@ -20,8 +20,9 @@ type AgentInstance struct { Model string Fallbacks []string Workspace string - MaxIterations int - ContextWindow int + MaxIterations int + TaskReminderInterval int + ContextWindow int Provider providers.LLMProvider Sessions *session.SessionManager ContextBuilder *ContextBuilder @@ -76,6 +77,11 @@ func NewAgentInstance( maxIter = 20 } + reminderInterval := defaults.TaskReminderInterval + if reminderInterval == 0 { + reminderInterval = 5 + } + // Resolve fallback candidates modelCfg := providers.ModelConfig{ Primary: model, @@ -89,8 +95,9 @@ func NewAgentInstance( Model: model, Fallbacks: fallbacks, Workspace: workspace, - MaxIterations: maxIter, - ContextWindow: defaults.MaxTokens, + MaxIterations: maxIter, + TaskReminderInterval: reminderInterval, + ContextWindow: defaults.MaxTokens, Provider: provider, Sessions: sessionsManager, ContextBuilder: contextBuilder, diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 6d0a61375..582eb9658 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -444,6 +444,40 @@ func (al *AgentLoop) runAgentLoop(ctx context.Context, agent *AgentInstance, opt return finalContent, nil } +// Task reminder constants and helpers. +const taskReminderMaxChars = 500 +const blockerMaxChars = 200 + +func shouldInjectReminder(iteration, interval int) bool { + if interval <= 0 { + return false + } + return iteration > 1 && iteration%interval == 0 +} + +func buildTaskReminder(userMessage string, lastBlocker string) providers.Message { + truncatedTask := utils.Truncate(userMessage, taskReminderMaxChars) + + var content string + if lastBlocker != "" { + truncatedBlocker := utils.Truncate(lastBlocker, blockerMaxChars) + content = fmt.Sprintf( + "[TASK REMINDER]\nOriginal task:\n---\n%s\n---\nLast blocker:\n---\n%s\n---\nDecide: fix the blocker if it's essential, or find an alternative approach to complete the original task.", + truncatedTask, truncatedBlocker, + ) + } else { + content = fmt.Sprintf( + "[TASK REMINDER]\nOriginal task:\n---\n%s\n---\nContinue with the next step toward completing this task.", + truncatedTask, + ) + } + + return providers.Message{ + Role: "user", + Content: content, + } +} + // runLLMIteration executes the LLM call loop with tool handling. func (al *AgentLoop) runLLMIteration(ctx context.Context, agent *AgentInstance, messages []providers.Message, opts processOptions) (string, int, error) { iteration := 0 @@ -611,6 +645,7 @@ func (al *AgentLoop) runLLMIteration(ctx context.Context, agent *AgentInstance, agent.Sessions.AddFullMessage(opts.SessionKey, assistantMsg) // Execute tool calls + var lastBlocker string for _, tc := range response.ToolCalls { argsJSON, _ := json.Marshal(tc.Arguments) argsPreview := utils.Truncate(string(argsJSON), 200) @@ -659,6 +694,11 @@ func (al *AgentLoop) runLLMIteration(ctx context.Context, agent *AgentInstance, contentForLLM = toolResult.Err.Error() } + // Track blockers for task reminder + if toolResult.IsError || toolResult.Err != nil { + lastBlocker = contentForLLM + } + toolResultMsg := providers.Message{ Role: "tool", Content: contentForLLM, @@ -669,6 +709,18 @@ func (al *AgentLoop) runLLMIteration(ctx context.Context, agent *AgentInstance, // Save tool result message to session agent.Sessions.AddFullMessage(opts.SessionKey, toolResultMsg) } + + // Inject ephemeral task reminder to prevent focus drift. + if shouldInjectReminder(iteration, agent.TaskReminderInterval) && !opts.NoHistory { + reminderMsg := buildTaskReminder(opts.UserMessage, lastBlocker) + messages = append(messages, reminderMsg) + logger.DebugCF("agent", "Injected task reminder", + map[string]interface{}{ + "agent_id": agent.ID, + "iteration": iteration, + "has_blocker": lastBlocker != "", + }) + } } return finalContent, iteration, nil diff --git a/pkg/agent/loop_test.go b/pkg/agent/loop_test.go index f2257973c..bd7b82386 100644 --- a/pkg/agent/loop_test.go +++ b/pkg/agent/loop_test.go @@ -5,6 +5,7 @@ import ( "fmt" "os" "path/filepath" + "strings" "testing" "time" @@ -628,3 +629,93 @@ func TestAgentLoop_ContextExhaustionRetry(t *testing.T) { t.Errorf("Expected history to be compressed (len < 8), got %d", len(finalHistory)) } } + +func TestShouldInjectReminder(t *testing.T) { + tests := []struct { + name string + iteration int + interval int + want bool + }{ + {"first iteration skipped", 1, 5, false}, + {"iteration 5 interval 5", 5, 5, true}, + {"iteration 10 interval 5", 10, 5, true}, + {"iteration 3 interval 5", 3, 5, false}, + {"interval zero disabled", 5, 0, false}, + {"interval negative disabled", 5, -1, false}, + {"iteration 2 interval 1", 2, 1, true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := shouldInjectReminder(tt.iteration, tt.interval) + if got != tt.want { + t.Errorf("shouldInjectReminder(%d, %d) = %v, want %v", tt.iteration, tt.interval, got, tt.want) + } + }) + } +} + +func TestBuildTaskReminder_WithoutBlocker(t *testing.T) { + msg := buildTaskReminder("implement feature X", "") + if msg.Role != "user" { + t.Errorf("expected role 'user', got %q", msg.Role) + } + if !strings.Contains(msg.Content, "[TASK REMINDER]") { + t.Error("expected content to contain '[TASK REMINDER]'") + } + if !strings.Contains(msg.Content, "implement feature X") { + t.Error("expected content to contain original message") + } + if strings.Contains(msg.Content, "blocker") { + t.Error("expected content NOT to contain 'blocker' when no blocker provided") + } + if !strings.Contains(msg.Content, "Continue with the next step") { + t.Error("expected content to contain continuation prompt") + } +} + +func TestBuildTaskReminder_WithBlocker(t *testing.T) { + msg := buildTaskReminder("implement feature X", "ModuleNotFoundError: No module named 'foo'") + if msg.Role != "user" { + t.Errorf("expected role 'user', got %q", msg.Role) + } + if !strings.Contains(msg.Content, "[TASK REMINDER]") { + t.Error("expected content to contain '[TASK REMINDER]'") + } + if !strings.Contains(msg.Content, "implement feature X") { + t.Error("expected content to contain original message") + } + if !strings.Contains(msg.Content, "Last blocker") { + t.Error("expected content to contain 'Last blocker'") + } + if !strings.Contains(msg.Content, "ModuleNotFoundError") { + t.Error("expected content to contain blocker text") + } +} + +func TestBuildTaskReminder_Truncation(t *testing.T) { + // Build a long message (1000 runes) + longMsg := strings.Repeat("あ", 1000) + longBlocker := strings.Repeat("X", 500) + + msg := buildTaskReminder(longMsg, longBlocker) + + // The full message should NOT contain 1000 'あ' characters + runeCount := strings.Count(msg.Content, "あ") + if runeCount >= 1000 { + t.Errorf("expected task message to be truncated, got %d 'あ' runes", runeCount) + } + // Should be at most taskReminderMaxChars (500) runes for the task part + if runeCount > taskReminderMaxChars { + t.Errorf("expected at most %d task runes, got %d", taskReminderMaxChars, runeCount) + } + + // Blocker should be truncated too + xCount := strings.Count(msg.Content, "X") + if xCount >= 500 { + t.Errorf("expected blocker to be truncated, got %d 'X' chars", xCount) + } + if xCount > blockerMaxChars { + t.Errorf("expected at most %d blocker chars, got %d", blockerMaxChars, xCount) + } +} diff --git a/pkg/config/config.go b/pkg/config/config.go index 682996bd6..d31dfc670 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -148,7 +148,8 @@ type AgentDefaults struct { ImageModelFallbacks []string `json:"image_model_fallbacks,omitempty"` MaxTokens int `json:"max_tokens" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_TOKENS"` Temperature float64 `json:"temperature" env:"PICOCLAW_AGENTS_DEFAULTS_TEMPERATURE"` - MaxToolIterations int `json:"max_tool_iterations" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_TOOL_ITERATIONS"` + MaxToolIterations int `json:"max_tool_iterations" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_TOOL_ITERATIONS"` + TaskReminderInterval int `json:"task_reminder_interval" env:"PICOCLAW_AGENTS_DEFAULTS_TASK_REMINDER_INTERVAL"` } type ChannelsConfig struct { @@ -331,7 +332,8 @@ func DefaultConfig() *Config { Model: "glm-4.7", MaxTokens: 8192, Temperature: 0.7, - MaxToolIterations: 20, + MaxToolIterations: 20, + TaskReminderInterval: 5, }, }, Channels: ChannelsConfig{