From b5b2ee43ec5cfdf2d9d78c6cbbac3b75d3e98971 Mon Sep 17 00:00:00 2001 From: Levi Neely <141506390+lneely@users.noreply.github.com> Date: Sun, 17 May 2026 10:14:44 +0200 Subject: [PATCH] loop: replace plan file re-injection with structured TaskState re-injection --- pkg/agent/core.go | 3 --- pkg/agent/loop.go | 46 +++++++++------------------------ pkg/agent/loop_maxsteps_test.go | 41 +++++++++++++++-------------- 3 files changed, 34 insertions(+), 56 deletions(-) diff --git a/pkg/agent/core.go b/pkg/agent/core.go index 340c60d..94cbcfd 100644 --- a/pkg/agent/core.go +++ b/pkg/agent/core.go @@ -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 diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 0b1498c..a8aea5f 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -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//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,18 +473,16 @@ 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 != "" { - 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, - ), - }, nil) + // 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: "[system: review your current task state and continue. Update it if your approach has changed.]\n\n" + msg, + }, nil) + } } } diff --git a/pkg/agent/loop_maxsteps_test.go b/pkg/agent/loop_maxsteps_test.go index f778fc3..79e1415 100644 --- a/pkg/agent/loop_maxsteps_test.go +++ b/pkg/agent/loop_maxsteps_test.go @@ -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 { @@ -155,27 +147,38 @@ func TestPlanReinjection(t *testing.T) { mb := &mockBackend{respond: sequentialStream(responses)} cfg := agentConfig{ - Backend: mb, - Tools: []backend.Tool{{Name: "tool"}}, - PlanFile: planPath, + Backend: mb, + Tools: []backend.Tool{{Name: "tool"}}, 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 }