perf: runtime task reminders

Author: voocel

internal/agent/runtime_policy.go

@@ -63,6 +63,10 @@ func (p *sessionRuntimePolicy) handleEvent(ev agentcore.Event) {
 		p.trackToolStart(ev)
 	case agentcore.EventToolExecEnd:
 		p.trackToolEnd(ev)
+	case agentcore.EventMessageEnd:
+		if msg, ok := ev.Message.(agentcore.Message); ok {
+			p.handleMessageEnd(msg)
+		}
 	}
 }
 
@@ -161,6 +165,25 @@ func (p *sessionRuntimePolicy) detectTaskManagementGap() {
 	}
 }
 
+func (p *sessionRuntimePolicy) handleMessageEnd(msg agentcore.Message) {
+	if msg.Role != agentcore.RoleAssistant {
+		return
+	}
+
+	s := p.session
+	if s.taskStore == nil {
+		return
+	}
+	snap := s.taskStore.Snapshot()
+	if key, reminder, ok := taskManagementReminderBeforeStop(msg, snap); ok {
+		s.deliverRuntimeReminder(
+			key,
+			ReminderTaskManagement,
+			reminder,
+		)
+	}
+}
+
 func hashToolArgs(raw json.RawMessage) string {
 	if len(raw) == 0 {
 		return ""

internal/agent/session_test.go

@@ -137,6 +137,78 @@ func (m *scriptedReminderModel) GenerateStream(
 
 func (m *scriptedReminderModel) SupportsTools() bool { return true }
 
+type taskCompletionReminderModel struct {
+	mu          sync.Mutex
+	callCount   int
+	sawReminder bool
+	taskID      string
+}
+
+func (m *taskCompletionReminderModel) Generate(
+	_ context.Context,
+	msgs []agentcore.Message,
+	_ []agentcore.ToolSpec,
+	_ ...agentcore.CallOption,
+) (*agentcore.LLMResponse, error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	sawInjectedReminder := false
+	for _, msg := range msgs {
+		if msg.Role == agentcore.RoleUser && strings.Contains(msg.TextContent(), "still in_progress") {
+			sawInjectedReminder = true
+			m.sawReminder = true
+		}
+	}
+
+	m.callCount++
+	switch m.callCount {
+	case 1:
+		return &agentcore.LLMResponse{Message: assistantTextMessage("总结完毕")}, nil
+	case 2:
+		if !sawInjectedReminder {
+			return &agentcore.LLMResponse{Message: assistantTextMessage("no reminder")}, nil
+		}
+		return &agentcore.LLMResponse{
+			Message: toolCallMessage(agentcore.ToolCall{
+				ID:   "task-update-1",
+				Name: "task_update",
+				Args: json.RawMessage(fmt.Sprintf(`{"taskId":%q,"status":"completed"}`, m.taskID)),
+			}),
+		}, nil
+	default:
+		return &agentcore.LLMResponse{Message: assistantTextMessage("已补上任务完成状态")}, nil
+	}
+}
+
+func (m *taskCompletionReminderModel) GenerateStream(
+	ctx context.Context,
+	msgs []agentcore.Message,
+	tools []agentcore.ToolSpec,
+	opts ...agentcore.CallOption,
+) (<-chan agentcore.StreamEvent, error) {
+	resp, err := m.Generate(ctx, msgs, tools, opts...)
+	if err != nil {
+		return nil, err
+	}
+	ch := make(chan agentcore.StreamEvent, 1)
+	ch <- agentcore.StreamEvent{
+		Type:       agentcore.StreamEventDone,
+		Message:    resp.Message,
+		StopReason: resp.Message.StopReason,
+	}
+	close(ch)
+	return ch, nil
+}
+
+func (m *taskCompletionReminderModel) SupportsTools() bool { return true }
+
+func (m *taskCompletionReminderModel) SawReminder() bool {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	return m.sawReminder
+}
+
 type namedChatModel struct {
 	name string
 }
@@ -856,6 +928,42 @@ func TestTaskManagementReminderQueuedForUntrackedOrBroadWork(t *testing.T) {
 	}
 }
 
+func TestTaskManagementReminderSteersBeforeStopWithOpenInProgressTask(t *testing.T) {
+	t.Parallel()
+
+	store := localtools.NewTaskStore()
+	task := store.Create("Summarize project state", "Write the final analysis summary", "Summarizing project state", nil)
+	inProgress := localtools.TaskInProgress
+	if _, err := store.Update(task.ID, localtools.TaskUpdateOpts{Status: &inProgress}); err != nil {
+		t.Fatalf("set task in_progress: %v", err)
+	}
+
+	model := &taskCompletionReminderModel{taskID: task.ID}
+	taskTools := localtools.NewTaskTools(store, nil, nil)
+	ag := agentcore.NewAgent(
+		agentcore.WithModel(model),
+		agentcore.WithTools(taskTools...),
+		agentcore.WithMaxTurns(10),
+	)
+	s := NewSession(SessionConfig{
+		Agent:     ag,
+		Settings:  config.Resolved{MaxTurns: 10},
+		Cwd:       t.TempDir(),
+		TaskStore: store,
+		Tools:     taskTools,
+	})
+	t.Cleanup(s.Close)
+
+	if err := s.Prompt("分析一下项目"); err != nil {
+		t.Fatalf("prompt: %v", err)
+	}
+
+	waitFor(t, time.Second, func() bool {
+		snap := store.Snapshot()
+		return model.SawReminder() && snap.Completed == 1 && snap.InProgress == 0 && s.LastAssistantText() == "已补上任务完成状态"
+	})
+}
+
 func TestRuntimeMetricsTrackCompactionSavings(t *testing.T) {
 	t.Parallel()
 

internal/agent/task_management_policy.go

@@ -1,14 +1,15 @@
 package agent
 
 import (
+	"fmt"
 	"strings"
 
 	"github.com/voocel/agentcore"
 	localtools "github.com/voocel/codebot/internal/tools"
 )
 
 const (
-	taskManagementPromptReminder       = "<system-reminder>\nThis is a multi-step implementation task. Create and maintain a task list before going deeper. Break the work into concrete tasks, mark a task in_progress before starting it, mark it completed immediately after finishing it, and keep moving to the next unblocked task.\n</system-reminder>"
+	taskManagementPromptReminder       = "<system-reminder>\nThis is a multi-step implementation task. Create and maintain a task list before going deeper. Break the work into concrete tasks, mark a task in_progress before starting it, mark it completed immediately after finishing it, and keep moving to the next unblocked task. Always mark the current task completed before giving your final answer unless the work is blocked, partial, or still failing verification.\n</system-reminder>"
 	taskManagementMissingReminder      = "<system-reminder>\nYou are doing multi-step work without maintaining a task list. Create concrete tasks now instead of continuing without structure.\n</system-reminder>"
 	taskManagementExpandSingleReminder = "<system-reminder>\nYour task list is too broad for the current scope. Split the single broad task into multiple more specific tasks and keep their statuses up to date.\n</system-reminder>"
 )
@@ -40,6 +41,31 @@ func taskManagementReminderForTurn(turn TurnOutcomeSnapshot, snap localtools.Tas
 	}
 }
 
+func taskManagementReminderBeforeStop(msg agentcore.Message, snap localtools.TaskSnapshot) (key, reminder string, ok bool) {
+	if msg.Role != agentcore.RoleAssistant || msg.StopReason != agentcore.StopReasonStop || snap.InProgress == 0 {
+		return "", "", false
+	}
+
+	inProgress := inProgressTasks(snap)
+	if len(inProgress) == 0 {
+		return "", "", false
+	}
+
+	taskRefs := make([]string, 0, len(inProgress))
+	taskIDs := make([]string, 0, len(inProgress))
+	for _, task := range inProgress {
+		taskIDs = append(taskIDs, task.ID)
+		taskRefs = append(taskRefs, fmt.Sprintf("#%s %s", task.ID, task.Subject))
+	}
+
+	return "task_management:before_stop_open:" + strings.Join(taskIDs, ","),
+		fmt.Sprintf(
+			"<system-reminder>\nYou are about to stop with task(s) still in_progress: %s. If this work is actually finished, call task_update to mark it completed before your final answer. If it is blocked, partial, or still failing verification, keep the task in_progress and explicitly say that instead of ending as if the work were done.\n</system-reminder>",
+			strings.Join(taskRefs, ", "),
+		),
+		true
+}
+
 func textContentFromBlocks(blocks []agentcore.ContentBlock) string {
 	var parts []string
 	for _, block := range blocks {
@@ -83,3 +109,17 @@ func looksLikeComplexTaskRequest(text string) bool {
 func allTasksCompleted(snap localtools.TaskSnapshot) bool {
 	return snap.Total > 0 && snap.Pending == 0 && snap.InProgress == 0
 }
+
+func inProgressTasks(snap localtools.TaskSnapshot) []localtools.Task {
+	if len(snap.Items) == 0 || snap.InProgress == 0 {
+		return nil
+	}
+
+	out := make([]localtools.Task, 0, snap.InProgress)
+	for _, task := range snap.Items {
+		if task.Status == localtools.TaskInProgress {
+			out = append(out, task)
+		}
+	}
+	return out
+}

internal/config/prompt.go

@@ -68,6 +68,7 @@ Break down and manage your work with task_create, task_update, and task_list.
 - Break larger requests into multiple specific tasks instead of one broad task
 - Mark a task in_progress before starting it
 - Mark a task completed immediately after fully finishing it
+- IMPORTANT: always mark the current task completed before giving your final answer, unless the work is blocked, partial, or still failing verification
 - Do not batch multiple completions together
 - Keep at most one task in_progress at a time
 - Check task_list before creating more tasks if a relevant task may already exist