loop: replace plan file re-injection with structured TaskState re-injection

This commit is contained in:
Levi Neely 2026-05-17 10:14:44 +02:00
parent 283a2e24d8
commit b5b2ee43ec
3 changed files with 34 additions and 56 deletions

View File

@ -364,9 +364,6 @@ func NewAgentCore(cfg AgentCoreConfig) Core {
}
if cfg.SessionID != "" {
os.MkdirAll(filepath.Join(ollieTmpDir(), cfg.SessionID), 0700) //nolint:errcheck
if mount := os.Getenv("OLLIE"); mount != "" {
run.PlanFile = filepath.Join(mount, "s", cfg.SessionID, "plan")
}
}
log := cfg.Log

View File

@ -1,12 +1,10 @@
package agent
import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"os"
"strings"
"sync"
"time"
@ -31,8 +29,7 @@ const (
)
// planReinjectInterval is the number of tool-call rounds between periodic
// plan re-injection. Every N steps the current contents of the plan file
// are surfaced back into the conversation to keep the model on track.
// task state re-injection into the conversation.
const planReinjectInterval = 10
type toolExecutor func(ctx context.Context, name string, args json.RawMessage) (string, error)
@ -56,27 +53,10 @@ type agentConfig struct {
// When reached, a soft nudge is injected and the loop exits cleanly.
// 0 means unlimited.
MaxSteps int
// PlanFile is the path to the session plan file (e.g. s/<id>/plan).
// When set, the loop re-surfaces its contents every planReinjectInterval
// tool-call rounds. Empty string disables plan re-injection.
PlanFile string
// IncrToolCallCount increments and returns the session-wide tool call counter.
IncrToolCallCount func() int64
}
// readPlanFile returns the contents of the plan file, or empty string if
// the file does not exist, is empty, or cannot be read.
func readPlanFile(path string) string {
if path == "" {
return ""
}
b, err := os.ReadFile(path)
if err != nil || len(bytes.TrimSpace(b)) == 0 {
return ""
}
return string(b)
}
func run(ctx context.Context, cfg agentConfig, state state) error {
var step int
var consecutiveErrors int // rounds where every tool call returned an error
@ -97,7 +77,7 @@ func run(ctx context.Context, cfg agentConfig, state state) error {
history = append([]backend.Message{{Role: "system", Content: cfg.preamble}}, history...)
}
if ts := state.taskState(); ts != nil {
if msg := ts.render(); msg != "" {
if msg := ts.render(); msg != "" && len(history) > 0 {
// Insert after system prompt, before conversation.
history = append([]backend.Message{history[0], {Role: "user", Content: msg}}, history[1:]...)
}
@ -493,20 +473,18 @@ func run(ctx context.Context, cfg agentConfig, state state) error {
break
}
// Periodically re-surface the plan file so it stays visible as tool
// results accumulate in the context. Only fires when PlanFile is set
// and the file has content.
if cfg.PlanFile != "" && step > 0 && step%planReinjectInterval == 0 {
if plan := readPlanFile(cfg.PlanFile); plan != "" {
// Periodic task state re-injection: every N tool rounds, surface the
// structured task state back into the conversation to keep the model on track.
if step > 0 && step%planReinjectInterval == 0 {
if ts := state.taskState(); ts != nil {
if msg := ts.render(); msg != "" {
state.update(backend.Message{
Role: "user",
Content: fmt.Sprintf(
"[system: your current plan:\n\n%s\n\nKeep this plan current — mark completed steps and revise if your approach has changed. Continue.]",
plan,
),
Content: "[system: review your current task state and continue. Update it if your approach has changed.]\n\n" + msg,
}, nil)
}
}
}
if cfg.AutoCompact != nil {
cfg.AutoCompact(ctx)

View File

@ -4,8 +4,6 @@ import (
"context"
"encoding/json"
"fmt"
"os"
"path/filepath"
"strings"
"testing"
@ -132,17 +130,11 @@ func TestMaxStepsExactBoundary(t *testing.T) {
}
}
// TestPlanReinjection verifies that the plan file contents are re-surfaced
// TestTaskStateReinjection verifies that the loop re-injects the task state
// into the conversation every planReinjectInterval tool rounds.
func TestPlanReinjection(t *testing.T) {
// Write a plan file to a temp location.
planPath := filepath.Join(t.TempDir(), "plan")
os.WriteFile(planPath, []byte("- [ ] step one\n- [ ] step two\n"), 0644)
func TestTaskStateReinjection(t *testing.T) {
// We need planReinjectInterval+1 tool rounds so the re-injection fires
// at step == planReinjectInterval (0-indexed, checked after increment).
// The check is: step > 0 && step%planReinjectInterval == 0, and step is
// incremented at the end of the loop, so it fires after the 10th round.
n := planReinjectInterval + 1
var responses []mockResponse
for i := range n {
@ -157,25 +149,36 @@ func TestPlanReinjection(t *testing.T) {
cfg := agentConfig{
Backend: mb,
Tools: []backend.Tool{{Name: "tool"}},
PlanFile: planPath,
Exec: func(_ context.Context, _ string, _ json.RawMessage) (string, error) {
return "ok", nil
},
}
s := newState()
s := &taskStateState{
simpleState: simpleState{},
ts: &TaskState{Objective: "test objective", PlanStep: "step one"},
}
if err := run(context.Background(), cfg, s); err != nil {
t.Fatalf("unexpected error: %v", err)
}
var planSeen bool
var seen bool
for _, m := range s.history() {
if m.Role == "user" && strings.Contains(m.Content, "step one") && strings.Contains(m.Content, "your current plan") {
planSeen = true
if m.Role == "user" && strings.Contains(m.Content, "test objective") && strings.Contains(m.Content, "task state") {
seen = true
break
}
}
if !planSeen {
t.Error("expected plan re-injection message in conversation history")
if !seen {
t.Error("expected task state re-injection message in conversation history")
}
}
// taskStateState wraps simpleState with a non-nil TaskState.
type taskStateState struct {
simpleState
ts *TaskState
}
func (s *taskStateState) taskState() *TaskState { return s.ts }
func (s *taskStateState) updateTaskState(ts TaskState) { s.ts = &ts }