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 != "" { if cfg.SessionID != "" {
os.MkdirAll(filepath.Join(ollieTmpDir(), cfg.SessionID), 0700) //nolint:errcheck 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 log := cfg.Log

View File

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

View File

@ -4,8 +4,6 @@ import (
"context" "context"
"encoding/json" "encoding/json"
"fmt" "fmt"
"os"
"path/filepath"
"strings" "strings"
"testing" "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. // into the conversation every planReinjectInterval tool rounds.
func TestPlanReinjection(t *testing.T) { func TestTaskStateReinjection(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)
// We need planReinjectInterval+1 tool rounds so the re-injection fires // We need planReinjectInterval+1 tool rounds so the re-injection fires
// at step == planReinjectInterval (0-indexed, checked after increment). // 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 n := planReinjectInterval + 1
var responses []mockResponse var responses []mockResponse
for i := range n { for i := range n {
@ -157,25 +149,36 @@ func TestPlanReinjection(t *testing.T) {
cfg := agentConfig{ cfg := agentConfig{
Backend: mb, Backend: mb,
Tools: []backend.Tool{{Name: "tool"}}, Tools: []backend.Tool{{Name: "tool"}},
PlanFile: planPath,
Exec: func(_ context.Context, _ string, _ json.RawMessage) (string, error) { Exec: func(_ context.Context, _ string, _ json.RawMessage) (string, error) {
return "ok", nil 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 { if err := run(context.Background(), cfg, s); err != nil {
t.Fatalf("unexpected error: %v", err) t.Fatalf("unexpected error: %v", err)
} }
var planSeen bool var seen bool
for _, m := range s.history() { for _, m := range s.history() {
if m.Role == "user" && strings.Contains(m.Content, "step one") && strings.Contains(m.Content, "your current plan") { if m.Role == "user" && strings.Contains(m.Content, "test objective") && strings.Contains(m.Content, "task state") {
planSeen = true seen = true
break break
} }
} }
if !planSeen { if !seen {
t.Error("expected plan re-injection message in conversation history") 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 }