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
This commit is contained in:
Levi Neely 2026-05-15 22:37:14 +02:00
parent 35a78169e5
commit 6dfae368d0
5 changed files with 128 additions and 21 deletions

View File

@ -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

View File

@ -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,
})
}
}

View File

@ -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")
}
}

View File

@ -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

View File

@ -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.