From 6dfae368d071311abe8a652fa969a51a5e5dc9b1 Mon Sep 17 00:00:00 2001
From: Levi Neely <141506390+lneely@users.noreply.github.com>
Date: Fri, 15 May 2026 22:37:14 +0200
Subject: [PATCH] fix: wire PlanFile, add missing rate-limit constants,
 simplify test helpers, fix regex

- Wire PlanFile from $OLLIE/s/{id}/plan into agentConfig so plan
  re-injection is no longer dead code
- Add missing failureWindow/maxFailures/blockDuration constants
- Replace hand-rolled contains() with strings.Contains
- Fix rm -rf .. regex to match end-of-string and whitespace
- Add TestPlanReinjection covering periodic plan re-surfacing
- Add loop_maxsteps_helpers_test.go with mockResponse/newState/sequentialStream
---
 pkg/agent/core.go                       |  3 +
 pkg/agent/loop_maxsteps_helpers_test.go | 60 +++++++++++++++++++
 pkg/agent/loop_maxsteps_test.go         | 78 ++++++++++++++++++-------
 pkg/tools/execute/code.go               |  2 +-
 pkg/tools/execute/server.go             |  6 ++
 5 files changed, 128 insertions(+), 21 deletions(-)
 create mode 100644 pkg/agent/loop_maxsteps_helpers_test.go

diff --git a/pkg/agent/core.go b/pkg/agent/core.go
index 04c51be..f1eb234 100644
--- a/pkg/agent/core.go
+++ b/pkg/agent/core.go
@@ -334,6 +334,9 @@ 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_maxsteps_helpers_test.go b/pkg/agent/loop_maxsteps_helpers_test.go
new file mode 100644
index 0000000..2e48bbd
--- /dev/null
+++ b/pkg/agent/loop_maxsteps_helpers_test.go
@@ -0,0 +1,60 @@
+package agent
+
+import (
+	"context"
+	"sync/atomic"
+
+	"ollie/pkg/backend"
+)
+
+// mockResponse defines a canned response for sequentialStream.
+type mockResponse struct {
+	content    string
+	toolCalls  []backend.ToolCall
+	stopReason string
+}
+
+// sequentialStream returns a respond function that plays back responses in order.
+func sequentialStream(responses []mockResponse) func(context.Context, []backend.Message, []backend.Tool, backend.GenerationParams) (<-chan backend.StreamEvent, error) {
+	var n int32
+	return func(_ context.Context, _ []backend.Message, _ []backend.Tool, _ backend.GenerationParams) (<-chan backend.StreamEvent, error) {
+		i := int(atomic.AddInt32(&n, 1)) - 1
+		if i >= len(responses) {
+			return textStream("done"), nil
+		}
+		r := responses[i]
+		ch := make(chan backend.StreamEvent, 1)
+		ch <- backend.StreamEvent{
+			Content:    r.content,
+			ToolCalls:  r.toolCalls,
+			Done:       true,
+			StopReason: r.stopReason,
+		}
+		close(ch)
+		return ch, nil
+	}
+}
+
+// simpleState is a minimal state implementation for direct run() calls.
+type simpleState struct {
+	msgs []backend.Message
+}
+
+func newState() *simpleState {
+	return &simpleState{}
+}
+
+func (s *simpleState) history() []backend.Message {
+	return s.msgs
+}
+
+func (s *simpleState) update(msg backend.Message, results []toolResult) {
+	s.msgs = append(s.msgs, msg)
+	for _, r := range results {
+		s.msgs = append(s.msgs, backend.Message{
+			Role:       "tool",
+			Content:    r.Content,
+			ToolCallID: r.ToolCallID,
+		})
+	}
+}
diff --git a/pkg/agent/loop_maxsteps_test.go b/pkg/agent/loop_maxsteps_test.go
index 1c0ef13..f778fc3 100644
--- a/pkg/agent/loop_maxsteps_test.go
+++ b/pkg/agent/loop_maxsteps_test.go
@@ -3,6 +3,10 @@ package agent
 import (
 	"context"
 	"encoding/json"
+	"fmt"
+	"os"
+	"path/filepath"
+	"strings"
 	"testing"
 
 	"ollie/pkg/backend"
@@ -13,11 +17,11 @@ import (
 func TestMaxStepsZeroUnlimited(t *testing.T) {
 	var steps int
 	mb := &mockBackend{
-		responses: []mockResponse{
+		respond: sequentialStream([]mockResponse{
 			{toolCalls: []backend.ToolCall{{ID: "1", Name: "tool", Arguments: json.RawMessage(`{}`)}}, stopReason: "tool_calls"},
 			{toolCalls: []backend.ToolCall{{ID: "2", Name: "tool", Arguments: json.RawMessage(`{}`)}}, stopReason: "tool_calls"},
 			{content: "done", stopReason: "stop"},
-		},
+		}),
 	}
 	cfg := agentConfig{
 		Backend:  mb,
@@ -45,11 +49,11 @@ func TestMaxStepsSoftNudge(t *testing.T) {
 
 	// Backend: two tool-calling rounds, then a final text turn.
 	mb := &mockBackend{
-		responses: []mockResponse{
+		respond: sequentialStream([]mockResponse{
 			{toolCalls: []backend.ToolCall{{ID: "1", Name: "tool", Arguments: json.RawMessage(`{}`)}}, stopReason: "tool_calls"},
 			{toolCalls: []backend.ToolCall{{ID: "2", Name: "tool", Arguments: json.RawMessage(`{}`)}}, stopReason: "tool_calls"},
 			{content: "wrapping up", stopReason: "stop"},
-		},
+		}),
 	}
 
 	// MaxSteps=1 means the guardrail fires after completing step 0 (the first
@@ -83,11 +87,9 @@ func TestMaxStepsSoftNudge(t *testing.T) {
 
 	// Confirm the nudge message is present in conversation history.
 	for _, m := range s.history() {
-		if m.Role == "user" && len(m.Content) > 0 {
-			if contains(m.Content, "step budget exhausted") {
-				nudgeSeen = true
-				break
-			}
+		if m.Role == "user" && strings.Contains(m.Content, "step budget exhausted") {
+			nudgeSeen = true
+			break
 		}
 	}
 	if !nudgeSeen {
@@ -112,7 +114,7 @@ func TestMaxStepsExactBoundary(t *testing.T) {
 	}
 	responses = append(responses, mockResponse{content: "done", stopReason: "stop"})
 
-	mb := &mockBackend{responses: responses}
+	mb := &mockBackend{respond: sequentialStream(responses)}
 	cfg := agentConfig{
 		Backend:  mb,
 		Tools:    []backend.Tool{{Name: "tool"}},
@@ -130,14 +132,50 @@ func TestMaxStepsExactBoundary(t *testing.T) {
 	}
 }
 
-func contains(s, substr string) bool {
-	return len(s) >= len(substr) && (s == substr || len(substr) == 0 ||
-		(func() bool {
-			for i := 0; i <= len(s)-len(substr); i++ {
-				if s[i:i+len(substr)] == substr {
-					return true
-				}
-			}
-			return false
-		})())
+// TestPlanReinjection verifies that the plan file contents are re-surfaced
+// 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)
+
+	// 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 {
+		responses = append(responses, mockResponse{
+			toolCalls:  []backend.ToolCall{{ID: fmt.Sprintf("%d", i+1), Name: "tool", Arguments: json.RawMessage(`{}`)}},
+			stopReason: "tool_calls",
+		})
+	}
+	responses = append(responses, mockResponse{content: "done", stopReason: "stop"})
+
+	mb := &mockBackend{respond: sequentialStream(responses)}
+	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()
+	if err := run(context.Background(), cfg, s); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	var planSeen 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
+			break
+		}
+	}
+	if !planSeen {
+		t.Error("expected plan re-injection message in conversation history")
+	}
 }
diff --git a/pkg/tools/execute/code.go b/pkg/tools/execute/code.go
index ecae941..d2c01fc 100644
--- a/pkg/tools/execute/code.go
+++ b/pkg/tools/execute/code.go
@@ -68,7 +68,7 @@ var bashPatterns = []*regexp.Regexp{
 	regexp.MustCompile(`rm\s+(-[a-z]*f[a-z]*\s+)*-[a-z]*r[a-z]*\s*/(home|var|usr|etc|boot|root|bin|sbin|lib|opt|srv)?`),
 	regexp.MustCompile(`rm\s+.*--recursive.*--force`),
 	regexp.MustCompile(`rm\s+.*--force.*--recursive`),
-	regexp.MustCompile(`rm\s+(-[a-z]*r[a-z]*\s+)*-[a-z]*f[a-z]*\s+\.\.[/]`),
+	regexp.MustCompile(`rm\s+(-[a-z]*r[a-z]*\s+)*-[a-z]*f[a-z]*\s+\.\.(/|\s|$)`),
 	regexp.MustCompile(`rm\s+(-[a-z]*r[a-z]*\s+)*-[a-z]*f[a-z]*\s+~`),
 	regexp.MustCompile(`rm\s+(-[a-z]*r[a-z]*\s+)*-[a-z]*f[a-z]*\s+\*`),
 	regexp.MustCompile(`:\s*\(\s*\)\s*\{\s*:\s*\|\s*:\s*&`), // fork bomb
diff --git a/pkg/tools/execute/server.go b/pkg/tools/execute/server.go
index 6fbbc50..fdfa2a8 100644
--- a/pkg/tools/execute/server.go
+++ b/pkg/tools/execute/server.go
@@ -20,6 +20,12 @@ import (
 	"ollie/pkg/tools"
 )
 
+const (
+	failureWindow = 1 * time.Minute
+	maxFailures   = 5
+	blockDuration = 5 * time.Minute
+)
+
 // Server runs code in a sandboxed environment.
 type Server struct {
 	// Confirm is an optional function called before executing sensitive operations.