| 1 | package agent |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "encoding/json" |
| 6 | "fmt" |
| 7 | "io" |
| 8 | "net/http" |
| 9 | "net/http/httptest" |
| 10 | "strings" |
| 11 | "testing" |
| 12 | |
| 13 | "reasonix/internal/event" |
| 14 | "reasonix/internal/tool" |
| 15 | ) |
| 16 | |
| 17 | // fatTool returns a fixed-size blob, standing in for a real read_file / bash |
| 18 | // whose output dominates the recent (verbatim-kept) tail of the session. |
| 19 | type fatTool struct{ blob string } |
| 20 | |
| 21 | func (fatTool) Name() string { return "fat_read" } |
| 22 | func (fatTool) Description() string { return "read a large file" } |
| 23 | func (fatTool) Schema() json.RawMessage { return json.RawMessage(`{"type":"object","properties":{}}`) } |
| 24 | func (fatTool) ReadOnly() bool { return true } |
| 25 | func (f fatTool) Execute(context.Context, json.RawMessage) (string, error) { |
| 26 | return f.blob, nil |
| 27 | } |
| 28 | |
| 29 | // loopMock emits exactly one tool call per user turn (a tool call when the last |
| 30 | // message is the user's, a final answer when it is the tool result), so each Run |
| 31 | // does one tool round — the step that triggers maybeCompact. finalText overrides |
| 32 | // the per-turn closing answer so a test can grow the session with assistant text |
| 33 | // (which pruning never touches) instead of tool output. |
| 34 | type loopMock struct { |
| 35 | t *testing.T |
| 36 | rounds int |
| 37 | finalText string |
| 38 | } |
| 39 | |
| 40 | func lastRole(msgs []json.RawMessage) string { |
| 41 | if len(msgs) == 0 { |
| 42 | return "" |
| 43 | } |
| 44 | var m struct { |
| 45 | Role string `json:"role"` |
| 46 | } |
| 47 | _ = json.Unmarshal(msgs[len(msgs)-1], &m) |
| 48 | return m.Role |
| 49 | } |
| 50 | |
| 51 | func (m *loopMock) handler(w http.ResponseWriter, r *http.Request) { |
| 52 | body, _ := io.ReadAll(r.Body) |
| 53 | if isSummarizeRequest(body) { |
| 54 | writeSSE(w, m.t, |
| 55 | streamChunk(deltaText("- goal: keep going\n- pending: continue the task")), |
| 56 | finishChunk("stop"), |
| 57 | usageChunk(80, 30, 0, 80)) |
| 58 | return |
| 59 | } |
| 60 | |
| 61 | msgs := decodeMessages(body) |
| 62 | promptTok := charsOf(msgs) / 4 |
| 63 | |
| 64 | if lastRole(msgs) == "tool" { |
| 65 | text := m.finalText |
| 66 | if text == "" { |
| 67 | text = "Done with this step." |
| 68 | } |
| 69 | writeSSE(w, m.t, |
| 70 | streamChunk(deltaText(text)), |
| 71 | finishChunk("stop"), |
| 72 | usageChunk(promptTok, 20, 0, promptTok)) |
| 73 | return |
| 74 | } |
| 75 | |
| 76 | m.rounds++ |
| 77 | writeSSE(w, m.t, |
| 78 | streamChunk(deltaToolCall(m.rounds, "fat_read", "{}")), |
| 79 | finishChunk("tool_calls"), |
| 80 | usageChunk(promptTok, 20, 0, promptTok)) |
| 81 | } |
| 82 | |
| 83 | // compactionsPerTurn drives `turns` user messages through a fresh agent wired to |
| 84 | // loopMock and reports, per turn, how many compactions started and whether an |
| 85 | // auto-compaction-paused notice was seen. |
| 86 | func compactionsPerTurn(t *testing.T, windowTok int, blob, finalText string, turns int) (perTurn []int, paused bool, prunes int) { |
| 87 | t.Helper() |
| 88 | mock := &loopMock{t: t, finalText: finalText} |
| 89 | srv := httptest.NewServer(http.HandlerFunc(mock.handler)) |
| 90 | defer srv.Close() |
| 91 | |
| 92 | reg := tool.NewRegistry() |
| 93 | reg.Add(fatTool{blob: blob}) |
| 94 | |
| 95 | a, _ := newAgent(t, srv.URL, reg, windowTok, 4) |
| 96 | started := 0 |
| 97 | a.sink = event.FuncSink(func(e event.Event) { |
| 98 | switch e.Kind { |
| 99 | case event.CompactionStarted: |
| 100 | started++ |
| 101 | case event.Notice: |
| 102 | if strings.Contains(e.Text, "Automatic context cleanup paused") { |
| 103 | paused = true |
| 104 | } |
| 105 | if strings.Contains(e.Text, "pruned") { |
| 106 | prunes++ |
| 107 | } |
| 108 | } |
| 109 | }) |
| 110 | |
| 111 | perTurn = make([]int, turns) |
| 112 | for i := 0; i < turns; i++ { |
| 113 | before := started |
| 114 | if err := a.Run(context.Background(), fmt.Sprintf("turn %d: keep going, continue the work", i)); err != nil { |
| 115 | t.Fatalf("Run %d: %v", i, err) |
| 116 | } |
| 117 | perTurn[i] = started - before |
| 118 | } |
| 119 | return perTurn, paused, prunes |
| 120 | } |
| 121 | |
| 122 | func consecutiveCompactingTurns(perTurn []int) int { |
| 123 | worst, run := 0, 0 |
| 124 | for _, n := range perTurn { |
| 125 | if n > 0 { |
| 126 | run++ |
| 127 | if run > worst { |
| 128 | worst = run |
| 129 | } |
| 130 | } else { |
| 131 | run = 0 |
| 132 | } |
| 133 | } |
| 134 | return worst |
| 135 | } |
| 136 | |
| 137 | // TestCompactionPausesWhenWindowTooSmall covers the user report: a tool output |
| 138 | // that alone exceeds the trigger used to make every "continue" turn re-compact |
| 139 | // forever. The stuck guard now caps it — at most two compactions, then a paused |
| 140 | // notice — instead of looping turn after turn. |
| 141 | func TestCompactionPausesWhenWindowTooSmall(t *testing.T) { |
| 142 | // One fat_read result (~1750 tok) exceeds the 0.8×1600 trigger on its own. |
| 143 | perTurn, paused, _ := compactionsPerTurn(t, 1600, strings.Repeat("LARGE FILE CONTENTS. ", 350), "", 8) |
| 144 | |
| 145 | total := 0 |
| 146 | for _, n := range perTurn { |
| 147 | total += n |
| 148 | } |
| 149 | t.Logf("compactions per turn: %v (total %d), paused=%v", perTurn, total, paused) |
| 150 | |
| 151 | if total > 2 { |
| 152 | t.Errorf("compacted %d times; the stuck guard should cap it at ≤2, not loop", total) |
| 153 | } |
| 154 | if !paused { |
| 155 | t.Errorf("expected an auto-compaction-paused notice") |
| 156 | } |
| 157 | } |
| 158 | |
| 159 | // TestCompactionHealthyWindowNeverLoops is the companion: when growth comes from |
| 160 | // assistant text (which pruning never touches), compaction still fires as the |
| 161 | // session grows but reclaims enough headroom that it never fires on consecutive |
| 162 | // turns and never trips the stuck guard. |
| 163 | func TestCompactionHealthyWindowNeverLoops(t *testing.T) { |
| 164 | perTurn, paused, _ := compactionsPerTurn(t, 40000, "small tool output", strings.Repeat("analysis paragraph. ", 600), 20) |
| 165 | |
| 166 | total := 0 |
| 167 | for _, n := range perTurn { |
| 168 | total += n |
| 169 | } |
| 170 | t.Logf("compactions per turn: %v (total %d), paused=%v", perTurn, total, paused) |
| 171 | |
| 172 | if paused { |
| 173 | t.Errorf("a healthy window should never pause auto-compaction") |
| 174 | } |
| 175 | if total == 0 { |
| 176 | t.Errorf("expected compaction to fire at least once over a long session") |
| 177 | } |
| 178 | if c := consecutiveCompactingTurns(perTurn); c > 1 { |
| 179 | t.Errorf("compaction fired on %d consecutive turns; a healthy compaction should leave breathing room", c) |
| 180 | } |
| 181 | } |
| 182 | |
| 183 | // TestPruneKeepsToolHeavySessionBounded: when growth comes from tool results, |
| 184 | // pruning alone keeps the prompt under the trigger for the whole session — the |
| 185 | // paid summarize call never happens and the stuck guard never trips. 20 turns of |
| 186 | // ~3k-token blobs would otherwise cross 0.8×40000 around turn 11. |
| 187 | func TestPruneKeepsToolHeavySessionBounded(t *testing.T) { |
| 188 | perTurn, paused, prunes := compactionsPerTurn(t, 40000, strings.Repeat("file line. ", 1100), "", 20) |
| 189 | |
| 190 | total := 0 |
| 191 | for _, n := range perTurn { |
| 192 | total += n |
| 193 | } |
| 194 | t.Logf("compactions per turn: %v (total %d), paused=%v, prunes=%d", perTurn, total, paused, prunes) |
| 195 | |
| 196 | if total != 0 { |
| 197 | t.Errorf("compaction fired %d times; pruning should keep a tool-heavy session bounded without folding", total) |
| 198 | } |
| 199 | if paused { |
| 200 | t.Errorf("auto-compaction paused; pruning should have prevented the stuck loop entirely") |
| 201 | } |
| 202 | if prunes == 0 { |
| 203 | t.Errorf("expected at least one prune pass over a tool-heavy session") |
| 204 | } |
| 205 | } |
| 206 |