diff --git a/pkg/agent/loop_prune.go b/pkg/agent/loop_prune.go index 429ea1e..a8332dd 100644 --- a/pkg/agent/loop_prune.go +++ b/pkg/agent/loop_prune.go @@ -35,7 +35,10 @@ func pruneStaleReads(history []backend.Message) []backend.Message { // recording new reads (a write within the same call supersedes // prior-call reads, not reads within this same call). for _, p := range writePaths { - for _, id := range reads[p] { + ids := reads[p] + // Keep the last (most recent) read as the base state; + // prune all earlier reads of this path. + for _, id := range ids[:max(0, len(ids)-1)] { stale[id] = true } delete(reads, p) diff --git a/pkg/agent/loop_prune_test.go b/pkg/agent/loop_prune_test.go index 50dd6de..d66bb44 100644 --- a/pkg/agent/loop_prune_test.go +++ b/pkg/agent/loop_prune_test.go @@ -2,6 +2,7 @@ package agent import ( "encoding/json" + "fmt" "testing" "ollie/pkg/backend" @@ -25,7 +26,8 @@ func makeToolResult(callID, content string) backend.Message { // TestPruneStaleReads_BasicWriteInvalidatesRead verifies that a file_read result // AND its corresponding assistant tool call are removed after a subsequent -// file_write to the same path. +// file_write to the same path. The last read before a write is kept as the +// base state; only earlier reads are pruned. func TestPruneStaleReads_BasicWriteInvalidatesRead(t *testing.T) { history := []backend.Message{ {Role: "user", Content: "edit foo"}, @@ -37,22 +39,9 @@ func TestPruneStaleReads_BasicWriteInvalidatesRead(t *testing.T) { out := pruneStaleReads(history) - // The file_read result (c1) and its assistant message should be gone. - for _, m := range out { - if m.Role == "tool" && m.ToolCallID == "c1" { - t.Error("stale file_read result was not pruned") - } - if m.Role == "assistant" { - for _, tc := range m.ToolCalls { - if tc.ID == "c1" { - t.Error("stale tool call c1 was not pruned from assistant message") - } - } - } - } - // user + assistant(c2) + tool(c2) = 3 messages - if len(out) != 3 { - t.Errorf("len(out) = %d; want 3", len(out)) + // c1 is the last (only) read before the write — kept as base state. + if len(out) != 5 { + t.Errorf("len(out) = %d; want 5 (single read before write is kept)", len(out)) } } @@ -107,8 +96,8 @@ func TestPruneStaleReads_NoWrites(t *testing.T) { } } -// TestPruneStaleReads_MultipleReadsOfSamePath prunes all stale reads of the -// same path, not just the most recent one. +// TestPruneStaleReads_MultipleReadsOfSamePath prunes all prior reads except +// the last one (which serves as the base state for the write). func TestPruneStaleReads_MultipleReadsOfSamePath(t *testing.T) { history := []backend.Message{ {Role: "user", Content: "x"}, @@ -122,20 +111,21 @@ func TestPruneStaleReads_MultipleReadsOfSamePath(t *testing.T) { out := pruneStaleReads(history) + // c1 (earlier read) should be pruned. for _, m := range out { - if m.Role == "tool" && (m.ToolCallID == "c1" || m.ToolCallID == "c2") { - t.Errorf("stale read %s was not pruned", m.ToolCallID) + if m.Role == "tool" && m.ToolCallID == "c1" { + t.Error("earlier read c1 should be pruned") } } - // c3 (write result) should still be present. - found := false + // c2 (last read before write) should be kept. + foundC2 := false for _, m := range out { - if m.Role == "tool" && m.ToolCallID == "c3" { - found = true + if m.Role == "tool" && m.ToolCallID == "c2" { + foundC2 = true } } - if !found { - t.Error("write result c3 should not have been pruned") + if !foundC2 { + t.Error("last read c2 should be kept as base state") } } @@ -152,10 +142,9 @@ func TestPruneStaleReads_FileEditAlsoPrunes(t *testing.T) { out := pruneStaleReads(history) - for _, m := range out { - if m.Role == "tool" && m.ToolCallID == "c1" { - t.Error("file_edit did not mark prior file_read as stale") - } + // c1 is the only (last) read before the edit — kept as base state. + if len(out) != 5 { + t.Errorf("len(out) = %d; want 5 (single read before edit is kept)", len(out)) } } @@ -178,8 +167,8 @@ func TestExtractExecCodePaths_Basic(t *testing.T) { } } -// TestPruneStaleReads_ReadWriteReadCycle is the primary correctness invariant: -// read before a write is pruned (both result and call); read after the same write is kept. +// TestPruneStaleReads_ReadWriteReadCycle: the last read before a write is kept +// as base state; a read after the write is also kept (it reflects current state). func TestPruneStaleReads_ReadWriteReadCycle(t *testing.T) { history := []backend.Message{ {Role: "user", Content: "update foo"}, @@ -193,36 +182,17 @@ func TestPruneStaleReads_ReadWriteReadCycle(t *testing.T) { out := pruneStaleReads(history) - // c1 (pre-write read) must be gone; c3 (post-write read) must remain. - for _, m := range out { - if m.Role == "tool" && m.ToolCallID == "c1" { - t.Error("pre-write read result c1 should have been pruned") - } - if m.Role == "assistant" { - for _, tc := range m.ToolCalls { - if tc.ID == "c1" { - t.Error("pre-write tool call c1 should have been pruned from assistant") - } - } - } - } - found := false - for _, m := range out { - if m.Role == "tool" && m.ToolCallID == "c3" { - found = true - } - } - if !found { - t.Error("post-write read result c3 should be kept") - } - // user + assistant(c2) + tool(c2) + assistant(c3) + tool(c3) = 5 messages - if len(out) != 5 { - t.Errorf("len(out) = %d; want 5", len(out)) + // c1 is the only read before the write — kept as base state. + // c3 is a read after the write — also kept. + // All messages should be preserved. + if len(out) != 7 { + t.Errorf("len(out) = %d; want 7 (both reads kept)", len(out)) } } // TestPruneStaleReads_AssistantWithTextContentKept verifies that an assistant -// message with text content is kept even when its tool calls are pruned. +// message with text content and a single read tool call is fully preserved +// when the read is the last one before a write (kept as base state). func TestPruneStaleReads_AssistantWithTextContentKept(t *testing.T) { // Build an assistant message with both text content and a tool call. args, _ := json.Marshal(map[string]any{"steps": []map[string]any{{"tool": "file_read", "args": []string{"/foo.go"}}}}) @@ -238,13 +208,17 @@ func TestPruneStaleReads_AssistantWithTextContentKept(t *testing.T) { out := pruneStaleReads(history) - // The assistant message with text should be kept (with tool calls removed). + // c1 is the last (only) read before the write — entire message preserved. + if len(out) != 5 { + t.Errorf("len(out) = %d; want 5 (all messages kept)", len(out)) + } + // The assistant message should retain both text and tool calls. var foundTextAssistant bool for _, m := range out { if m.Role == "assistant" && m.Content == "Let me read that file." { foundTextAssistant = true - if len(m.ToolCalls) != 0 { - t.Error("stale tool call should have been removed from assistant message") + if len(m.ToolCalls) != 1 { + t.Error("tool call should be preserved (read is kept as base state)") } } } @@ -253,6 +227,128 @@ func TestPruneStaleReads_AssistantWithTextContentKept(t *testing.T) { } } +// TestPruneStaleReads_ReasoningPreservesEvidence verifies that when the model +// reads a file, reasons about it, then edits it, the read is kept as base +// state — preserving the evidence that grounds the reasoning. +func TestPruneStaleReads_ReasoningPreservesEvidence(t *testing.T) { + args1, _ := json.Marshal(map[string]any{"steps": []map[string]any{{"tool": "file_read", "args": []string{"/foo.go"}}}}) + history := []backend.Message{ + {Role: "user", Content: "fix the bug in foo.go"}, + {Role: "assistant", Content: "Let me read the file to understand the issue.", ToolCalls: []backend.ToolCall{ + {ID: "c1", Name: "execute_code", Arguments: json.RawMessage(args1)}, + }}, + makeToolResult("c1", "func handler() {\n\treturn nil // BUG: should return error\n}"), + {Role: "assistant", Content: "I can see the bug on line 2: the handler returns nil instead of an error. I'll fix this."}, + {Role: "user", Content: "yes, fix it"}, + makeExecCodeMsg("c2", []map[string]any{{"tool": "file_edit", "args": []string{"/foo.go", "return nil", "return fmt.Errorf(...)"}}}), + makeToolResult("c2", "ok"), + } + + out := pruneStaleReads(history) + + // c1 is the last (only) read before the edit — kept as base state. + // The reasoning and all other messages are also preserved. + if len(out) != 7 { + t.Errorf("len(out) = %d; want 7 (all messages preserved)", len(out)) + } +} + +// TestPruneStaleReads_ReadEditReadEditChain shows that in a read-edit-read-edit +// cycle, each read is the last read before its respective edit and is kept. +func TestPruneStaleReads_ReadEditReadEditChain(t *testing.T) { + history := []backend.Message{ + {Role: "user", Content: "update foo.go twice"}, + // First read + makeExecCodeMsg("c1", []map[string]any{{"tool": "file_read", "args": []string{"/foo.go"}}}), + makeToolResult("c1", "version 1"), + // First edit + makeExecCodeMsg("c2", []map[string]any{{"tool": "file_edit", "args": []string{"/foo.go", "old", "new"}}}), + makeToolResult("c2", "ok"), + // Verification read + makeExecCodeMsg("c3", []map[string]any{{"tool": "file_read", "args": []string{"/foo.go"}}}), + makeToolResult("c3", "version 2"), + // Second edit + makeExecCodeMsg("c4", []map[string]any{{"tool": "file_edit", "args": []string{"/foo.go", "new", "newer"}}}), + makeToolResult("c4", "ok"), + } + + out := pruneStaleReads(history) + + // c1 is the last read before c2 — kept as base for first edit. + // c3 is the last read before c4 — kept as base for second edit. + // All messages preserved. + if len(out) != 9 { + t.Errorf("len(out) = %d; want 9 (both reads kept as base states)", len(out)) + } +} + +// TestPruneStaleReads_MultiFileTurnAccumulation shows that in a multi-file +// editing session, each file's last read is kept as base state. +func TestPruneStaleReads_MultiFileTurnAccumulation(t *testing.T) { + history := []backend.Message{ + {Role: "user", Content: "refactor: rename getUserName to getUsername across the codebase"}, + } + + // Simulate: read 5 files, then edit 3 of them. + files := []string{"/a.go", "/b.go", "/c.go", "/d.go", "/e.go"} + for i, f := range files { + id := fmt.Sprintf("r%d", i) + history = append(history, + makeExecCodeMsg(id, []map[string]any{{"tool": "file_read", "args": []string{f}}}), + makeToolResult(id, fmt.Sprintf("contents of %s with getUserName", f)), + ) + } + // Edit a.go, c.go, e.go + for i, f := range []string{"/a.go", "/c.go", "/e.go"} { + id := fmt.Sprintf("w%d", i) + history = append(history, + makeExecCodeMsg(id, []map[string]any{{"tool": "file_edit", "args": []string{f, "getUserName", "getUsername"}}}), + makeToolResult(id, "ok"), + ) + } + + out := pruneStaleReads(history) + + // Each edited file has exactly one read — that's the last read, so it's kept. + // All 5 reads are kept (each is the only/last read of its path). + kept := 0 + for _, m := range out { + if m.Role == "tool" && len(m.ToolCallID) == 2 && m.ToolCallID[0] == 'r' { + kept++ + } + } + if kept != 5 { + t.Errorf("kept %d reads; want 5 (all are last reads of their paths)", kept) + } + + // Total: 1 user + 5*(read assistant + read tool) + 3*(edit assistant + edit tool) = 1 + 10 + 6 = 17 + if len(out) != 17 { + t.Errorf("len(out) = %d; want 17", len(out)) + } +} + +// TestPruneStaleReads_ParallelReadWrite tests a single execute_code call +// containing both a file_read and file_write to the same path in parallel steps. +func TestPruneStaleReads_ParallelReadWrite(t *testing.T) { + // A single call that reads and writes the same file. + history := []backend.Message{ + {Role: "user", Content: "go"}, + makeExecCodeMsg("c1", []map[string]any{ + {"tool": "file_read", "args": []string{"/foo.go"}}, + {"tool": "file_write", "args": []string{"/foo.go", "new"}}, + }), + makeToolResult("c1", "read result\nwrite ok"), + } + + out := pruneStaleReads(history) + + // Per the code comment: "a write within the same call supersedes prior-call + // reads, not reads within this same call." So c1 should NOT be pruned. + if len(out) != len(history) { + t.Errorf("len(out) = %d; want %d (same-call read+write should not self-prune)", len(out), len(history)) + } +} + // TestPruneStaleReads_NonExecuteCodeUnaffected verifies that tool calls with // names other than "execute_code" are ignored by the path tracker and never // cause unrelated tool results to be pruned.