lkn
/
ollie-core
Archived
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

pruneStaleReads: keep last read per path as base state 

Previously, all reads of a path were pruned when a write occurred. This
caused the model to lose the base state of files it edited, leaving its
reasoning ungrounded and edit diffs without context.

Now only earlier reads are pruned; the most recent read before a write is
kept. Combined with the edit diff in the tool result, the model retains
full knowledge: base contents + deltas = current state.
This commit is contained in:
Levi Neely 2026-04-29 18:03:49 +02:00
parent ed3373d3f0
commit 2d81cc2aa1
2 changed files with 162 additions and 63 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
pkg/agent/loop_prune.go
View File

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

220
pkg/agent/loop_prune_test.go
View File

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