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:
parent
35a78169e5
commit
6dfae368d0
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
@ -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")
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
Reference in New Issue