diff --git a/docs/changes/unreleased/2045-opened-chat-memory.md b/docs/changes/unreleased/2045-opened-chat-memory.md new file mode 100644 index 0000000000..3f89a1d9da --- /dev/null +++ b/docs/changes/unreleased/2045-opened-chat-memory.md @@ -0,0 +1,16 @@ +--- +kind: fixed +title: Opening a long chat no longer holds its whole journal in memory +pr: 2045 +surface: [engine] +invalidates: + - "Previously an opened chat kept the words of every journal line its replay read, including the copies a later compaction re-recorded; now it keeps only the messages it holds." +--- + +A resumed chat now lets go of every message line a later compaction replaced, +at the open and at each compaction, rewind or fold while it runs, as it +already did for the messages' presentation marks. A journal too large to keep +in the reading cache is read into a large amount of short-lived memory; once +such a reading is done, the engine returns that memory to the system in the +background. Opening a 1,371-entry chat on a Mac went from about 720 MB to +about 80 MB for the engine. Entries, their pages and times are unchanged. diff --git a/internal/session/agent.go b/internal/session/agent.go index c34a02ebce..e9959c8a7a 100644 --- a/internal/session/agent.go +++ b/internal/session/agent.go @@ -392,6 +392,7 @@ func newAgent(config Config, client Completer) (*Agent, error) { // The region above the last pass is shaped and let go of, so its marks // go with it ([presentationIndex.retain]). agent.presentation.retain(agent.messages) + file.retainStamps(agent.messages) agent.restorePasses(replayed.passes) // AND THE PICTURES ARE MADE SAFE HERE RATHER THAN IN THE REPLAY. A // session resumed onto a model without vision — the journal remembers @@ -423,6 +424,12 @@ func newAgent(config Config, client Completer) (*Agent, error) { // that has not spoken since the field arrived would read as free on // home forever. agent.stampRestoredSpend(agent.usage) + // The earlier region has been shaped and its message indexes pruned. + // Let the replay go before returning a large read's working memory. + readBytes := replayed.readBytes + replayed = replayedSession{} + restored = nil + releaseAfterLargeRead(readBytes) } // AND WHERE THIS CONVERSATION IS ABOUT. The referred places ride on the same // meta.json everything else about the session's identity does (places.go), so @@ -3170,9 +3177,10 @@ func (a *Agent) rememberPersonLocked(user userMessage, kept ai.Message) func() { return func() { a.presentation.forget(kept) } } -// retainPresentationLocked cuts the presentation marks back to the messages -// this session still holds — the transcript and every queue of messages not yet -// in it — after a boundary that let messages go ([presentationIndex.retain]). +// retainPresentationLocked cuts the presentation marks and journal stamps back +// to the messages this session still holds — the transcript and every queue of +// messages not yet in it — after a boundary that let messages go +// ([presentationIndex.retain]). func (a *Agent) retainPresentationLocked() { queued := make([]ai.Message, 0, len(a.steering)+len(a.ambient)+len(a.followups)) for _, user := range a.steering { @@ -3185,6 +3193,7 @@ func (a *Agent) retainPresentationLocked() { queued = append(queued, follow.message.message) } a.presentation.retain(a.messages, queued) + a.file.retainStamps(a.messages, queued) } func (a *Agent) record(message ai.Message) { diff --git a/internal/session/entrytime.go b/internal/session/entrytime.go index 65e4fcec13..b7f730b92f 100644 --- a/internal/session/entrytime.go +++ b/internal/session/entrytime.go @@ -61,6 +61,58 @@ func (s *stampIndex) of(message ai.Message) time.Time { return s.at[key] } +// retain keeps only the instants of messages something still holds. Like a +// presentation mark, a stamp's key holds the message's whole part array alive; +// keeping stamps for folded or rewound messages would keep their words too. +func (s *stampIndex) retain(live ...[]ai.Message) { + if s == nil { + return + } + s.mu.Lock() + defer s.mu.Unlock() + if len(s.at) == 0 { + return + } + kept := make(map[any]time.Time) + for _, messages := range live { + for index := range messages { + if key := stampKey(messages[index]); key != nil { + if at, ok := s.at[key]; ok { + kept[key] = at + } + } + } + } + s.at = kept +} + +// retainStamps follows the live transcript and its queues at the same boundaries +// as the presentation index, after any departing region has been shaped. +func (s *sessionFile) retainStamps(live ...[]ai.Message) { + if s == nil { + return + } + s.mu.Lock() + defer s.mu.Unlock() + s.stamps.retain(live...) +} + +// rememberStamp puts back the instant a message was written, for a message a +// boundary let go of and an undo is handing back ([Agent.UndoRewind]): the +// stamps field is read under the file's lock everywhere else. +func (s *sessionFile) rememberStamp(message ai.Message, at time.Time) { + if s == nil { + return + } + s.mu.Lock() + if s.stamps == nil { + s.stamps = &stampIndex{} + } + stamps := s.stamps + s.mu.Unlock() + stamps.remember(message, at) +} + // lineInstant reads a journal line's own timestamp, and the zero time for a // line that has none or one this build cannot read. func lineInstant(text string) time.Time { diff --git a/internal/session/readrelease.go b/internal/session/readrelease.go new file mode 100644 index 0000000000..be6383f10d --- /dev/null +++ b/internal/session/readrelease.go @@ -0,0 +1,74 @@ +package session + +import ( + "runtime" + "sync" + "time" +) + +// releaseCooldown keeps sustained journal paging from repeatedly paying for a +// process-wide collection while the runtime scavenger returns the freed heap. +const releaseCooldown = 5 * time.Second + +// freeJournalReadMemory is replaceable under journalReadRelease.mu so tests can +// hold a collection in progress. GC lowers the heap goal after a large read, +// letting the runtime's paced background scavenger return memory without the +// process-wide forced scavenge that debug.FreeOSMemory would add on every run. +var freeJournalReadMemory = runtime.GC + +// The clock and scheduler are replaced under journalReadRelease.mu so cooldown +// tests can advance time without waiting for real collections or timers. +var journalReadNow = time.Now +var journalReadAfterFunc = func(delay time.Duration, run func()) { time.AfterFunc(delay, run) } + +var journalReadRelease struct { + mu sync.Mutex + running bool + pending bool + scheduled bool + lastStart time.Time +} + +// releaseAfterLargeRead never waits for collection. Requests share one trailing +// collection at the cooldown's end, even while an earlier collection is running. +func releaseAfterLargeRead(size int64) { + if size < recordCacheFileBytes { + return + } + journalReadRelease.mu.Lock() + defer journalReadRelease.mu.Unlock() + journalReadRelease.pending = true + scheduleJournalReadReleaseLocked() +} + +func scheduleJournalReadReleaseLocked() { + if journalReadRelease.running || journalReadRelease.scheduled || !journalReadRelease.pending { + return + } + now := journalReadNow() + if delay := journalReadRelease.lastStart.Add(releaseCooldown).Sub(now); !journalReadRelease.lastStart.IsZero() && delay > 0 { + journalReadRelease.scheduled = true + journalReadAfterFunc(delay, func() { + journalReadRelease.mu.Lock() + defer journalReadRelease.mu.Unlock() + journalReadRelease.scheduled = false + scheduleJournalReadReleaseLocked() + }) + return + } + journalReadRelease.pending = false + journalReadRelease.running = true + free := freeJournalReadMemory + go func() { + journalReadRelease.mu.Lock() + journalReadRelease.lastStart = journalReadNow() + journalReadRelease.mu.Unlock() + free() + journalReadRelease.mu.Lock() + defer journalReadRelease.mu.Unlock() + journalReadRelease.running = false + // A request during GC keeps the deadline measured from its start. If GC + // outlasts that deadline, the pending run can start as soon as it finishes. + scheduleJournalReadReleaseLocked() + }() +} diff --git a/internal/session/readrelease_test.go b/internal/session/readrelease_test.go new file mode 100644 index 0000000000..757cb9b20e --- /dev/null +++ b/internal/session/readrelease_test.go @@ -0,0 +1,303 @@ +package session + +import ( + "os" + "path/filepath" + "strings" + "sync" + "sync/atomic" + "testing" + "time" +) + +func waitReadReleaseRunFinished(t *testing.T) { + t.Helper() + deadline := time.Now().Add(5 * time.Second) + for { + journalReadRelease.mu.Lock() + running := journalReadRelease.running + journalReadRelease.mu.Unlock() + if !running { + return + } + if time.Now().After(deadline) { + t.Fatal("journal read collection did not finish") + } + time.Sleep(time.Millisecond) + } +} + +type readReleaseTimer struct { + at time.Time + run func() +} + +type readReleaseClock struct { + mu sync.Mutex + now time.Time + timer *readReleaseTimer + schedules int +} + +func (c *readReleaseClock) instant() time.Time { + c.mu.Lock() + defer c.mu.Unlock() + return c.now +} + +func (c *readReleaseClock) afterFunc(delay time.Duration, run func()) { + c.mu.Lock() + defer c.mu.Unlock() + if c.timer != nil { + panic("more than one journal read release timer") + } + c.schedules++ + c.timer = &readReleaseTimer{at: c.now.Add(delay), run: run} +} + +func (c *readReleaseClock) advance(d time.Duration) { + c.mu.Lock() + target := c.now.Add(d) + for c.timer != nil && !c.timer.at.After(target) { + timer := c.timer + c.now, c.timer = timer.at, nil + c.mu.Unlock() + timer.run() + c.mu.Lock() + } + c.now = target + c.mu.Unlock() +} + +func replaceReadRelease(t *testing.T, free func()) *readReleaseClock { + t.Helper() + waitReadReleaseRunFinished(t) + clock := &readReleaseClock{now: time.Date(2026, 10, 15, 14, 0, 0, 0, time.UTC)} + journalReadRelease.mu.Lock() + priorFree, priorNow, priorAfter := freeJournalReadMemory, journalReadNow, journalReadAfterFunc + if journalReadRelease.scheduled || journalReadRelease.pending { + journalReadRelease.mu.Unlock() + t.Fatal("prior journal read release still scheduled") + } + priorStart := journalReadRelease.lastStart + journalReadRelease.lastStart = time.Time{} + freeJournalReadMemory, journalReadNow, journalReadAfterFunc = free, clock.instant, clock.afterFunc + journalReadRelease.mu.Unlock() + t.Cleanup(func() { + waitReadReleaseRunFinished(t) + journalReadRelease.mu.Lock() + journalReadRelease.pending, journalReadRelease.scheduled = false, false + journalReadRelease.lastStart = priorStart + freeJournalReadMemory, journalReadNow, journalReadAfterFunc = priorFree, priorNow, priorAfter + journalReadRelease.mu.Unlock() + }) + return clock +} + +func receiveReadRelease(t *testing.T, ch <-chan struct{}) { + t.Helper() + select { + case <-ch: + case <-time.After(5 * time.Second): + t.Fatal("journal read collection did not start or caller blocked") + } +} + +func TestLargeReadReleaseThreshold(t *testing.T) { + var calls atomic.Int32 + clock := replaceReadRelease(t, func() { calls.Add(1) }) + releaseAfterLargeRead(recordCacheFileBytes - 1) + waitReadReleaseRunFinished(t) + if n := calls.Load(); n != 0 { + t.Fatalf("small read caused %d collections", n) + } + releaseAfterLargeRead(recordCacheFileBytes) + waitReadReleaseRunFinished(t) + clock.advance(2 * releaseCooldown) + waitReadReleaseRunFinished(t) + if n := calls.Load(); n != 1 || clock.schedules != 0 { + t.Fatalf("one large read caused %d collections and %d timers, want one prompt collection", n, clock.schedules) + } +} + +func TestLargeReadReleaseBurstHasOneTrailingCollection(t *testing.T) { + var calls atomic.Int32 + clock := replaceReadRelease(t, func() { calls.Add(1) }) + releaseAfterLargeRead(recordCacheFileBytes) + waitReadReleaseRunFinished(t) + for range 64 { + releaseAfterLargeRead(recordCacheFileBytes) + } + clock.advance(releaseCooldown - time.Nanosecond) + waitReadReleaseRunFinished(t) + if n := calls.Load(); n != 1 || clock.schedules != 1 { + t.Fatalf("burst before deadline: %d collections, %d timers, want 1 each", n, clock.schedules) + } + clock.advance(time.Nanosecond) + waitReadReleaseRunFinished(t) + clock.advance(2 * releaseCooldown) + waitReadReleaseRunFinished(t) + if n := calls.Load(); n != 2 || clock.schedules != 1 { + t.Fatalf("burst: %d collections, %d timers, want one prompt and one trailing collection", n, clock.schedules) + } +} + +func TestLargeReadReleaseSustainedRequestsStayBounded(t *testing.T) { + var startsMu sync.Mutex + var starts []time.Time + clock := replaceReadRelease(t, func() { + journalReadRelease.mu.Lock() + at := journalReadRelease.lastStart + journalReadRelease.mu.Unlock() + startsMu.Lock() + starts = append(starts, at) + startsMu.Unlock() + }) + start := clock.instant() + releaseAfterLargeRead(recordCacheFileBytes) + waitReadReleaseRunFinished(t) + const duration, interval = 30 * time.Second, 900 * time.Millisecond + for elapsed := interval; elapsed < duration; elapsed += interval { + clock.advance(interval) + waitReadReleaseRunFinished(t) + releaseAfterLargeRead(recordCacheFileBytes) + } + lastRequest := clock.instant() + clock.advance(releaseCooldown) + waitReadReleaseRunFinished(t) + startsMu.Lock() + defer startsMu.Unlock() + limit := int((duration+releaseCooldown-1)/releaseCooldown) + 1 + if len(starts) > limit || len(starts) < 2 || !starts[len(starts)-1].After(lastRequest) { + t.Fatalf("sustained requests: %d collections, limit %d, starts %v, last request %v", len(starts), limit, starts, lastRequest) + } + for i := 1; i < len(starts); i++ { + if gap := starts[i].Sub(starts[i-1]); gap < releaseCooldown { + t.Fatalf("collections %d and %d separated by %v, want at least %v", i-1, i, gap, releaseCooldown) + } + } + if !starts[0].Equal(start) { + t.Fatalf("first collection delayed until %v, want %v", starts[0], start) + } + t.Logf("30 s of sustained requests: %d collections (limit %d), final collection after last request", len(starts), limit) +} + +func TestLargeReadReleaseCoalescesWithoutBlocking(t *testing.T) { + for _, blocked := range []time.Duration{time.Second, 2 * releaseCooldown} { + t.Run(blocked.String(), func(t *testing.T) { + entered := make(chan struct{}, 4) + resume := make(chan struct{}, 2) + var calls atomic.Int32 + clock := replaceReadRelease(t, func() { + calls.Add(1) + entered <- struct{}{} + <-resume + }) + // Unblock a failed assertion before restoring the hook in the prior cleanup. + t.Cleanup(func() { close(resume) }) + returned := make(chan struct{}) + go func() { + releaseAfterLargeRead(recordCacheFileBytes) + close(returned) + }() + receiveReadRelease(t, entered) + receiveReadRelease(t, returned) + var requests sync.WaitGroup + for range 64 { + requests.Go(func() { releaseAfterLargeRead(recordCacheFileBytes) }) + } + done := make(chan struct{}) + go func() { requests.Wait(); close(done) }() + receiveReadRelease(t, done) + clock.advance(blocked) + if n := calls.Load(); n != 1 { + t.Fatalf("collection overlapped a blocked run: %d calls", n) + } + resume <- struct{}{} + if blocked < releaseCooldown { + waitReadReleaseRunFinished(t) + clock.advance(releaseCooldown - blocked - time.Nanosecond) + if n := calls.Load(); n != 1 { + t.Fatalf("follow-up started before cooldown: %d calls", n) + } + clock.advance(time.Nanosecond) + } + receiveReadRelease(t, entered) + resume <- struct{}{} + waitReadReleaseRunFinished(t) + clock.advance(2 * releaseCooldown) + waitReadReleaseRunFinished(t) + if n := calls.Load(); n != 2 { + t.Fatalf("concurrent requests caused %d collections, want one plus one follow-up", n) + } + }) + } +} + +// Whitespace lines grow the journal without making the fixture retain a large +// display record or depending on the scanner's limit for one very long line. +func readReleaseJournal(t *testing.T, size int64) string { + t.Helper() + path := filepath.Join(t.TempDir(), "session.jsonl") + f, err := os.Create(path) + if err != nil { + t.Fatal(err) + } + defer f.Close() + header := `{"type":"session","version":1,"id":"read-release"} +{"type":"message","role":"user","content":"kept words","timestamp":"2026-10-15T14:00:00Z"} +{"type":"compaction","folded":1,"window":1,"timestamp":"2026-10-15T14:01:00Z"} +{"type":"message","role":"user","content":"kept words","timestamp":"2026-10-15T14:01:00Z"} +` + if _, err := f.WriteString(header); err != nil { + t.Fatal(err) + } + padding := strings.Repeat(" ", (64<<10)-1) + "\n" + for left := size - int64(len(header)); left > 0; { + n := min(left, int64(len(padding))) + if _, err := f.WriteString(padding[len(padding)-int(n):]); err != nil { + t.Fatal(err) + } + left -= n + } + return path +} + +func TestRecordReadReleaseFollowsCallback(t *testing.T) { + for name, size := range map[string]int64{"below threshold": recordCacheFileBytes - 1, "at threshold": recordCacheFileBytes} { + t.Run(name, func(t *testing.T) { + path := readReleaseJournal(t, size) + var calls atomic.Int32 + replaceReadRelease(t, func() { calls.Add(1) }) + // Both reads are refused retention, even the one just under the + // threshold; only the shared large-file threshold pays for a GC. + cache := settledCache(RecordCacheLimits{FileBytes: 1}) + if err := cache.WithRecord(path, func(Record) error { + if calls.Load() != 0 { + t.Fatal("collection ran before the callback returned") + } + return nil + }); err != nil { + t.Fatal(err) + } + waitReadReleaseRunFinished(t) + want := int32(0) + if size >= recordCacheFileBytes { + want = 1 + } + if n := calls.Load(); n != want { + t.Fatalf("collections = %d, want %d", n, want) + } + }) + } +} + +func TestResumeReadReleaseAfterStampPrune(t *testing.T) { + path := readReleaseJournal(t, recordCacheFileBytes) + called := make(chan struct{}, 1) + replaceReadRelease(t, func() { called <- struct{}{} }) + a, _ := newTestAgent(t, &refusingCompleter{t}, func(c *Config) { c.SessionFile = path }) + receiveReadRelease(t, called) + waitReadReleaseRunFinished(t) + assertLiveStamps(t, a, journalStampInstants(t, a, nil)) +} diff --git a/internal/session/recordcache.go b/internal/session/recordcache.go index e9d3256e19..3119019237 100644 --- a/internal/session/recordcache.go +++ b/internal/session/recordcache.go @@ -389,6 +389,7 @@ func (c *RecordCache) parseLocked(path string, size int64, stamp recordStamp, st c.mu.Lock() c.reserved -= cost c.mu.Unlock() + releaseAfterLargeRead(size) }, nil } diff --git a/internal/session/rewind.go b/internal/session/rewind.go index 00e019737a..dcd8f2fad7 100644 --- a/internal/session/rewind.go +++ b/internal/session/rewind.go @@ -8,6 +8,7 @@ import ( "errors" "fmt" "strings" + "time" "github.com/Agent-Field/agentfield/sdk/go/ai" "github.com/Agent-Field/codeaf/internal/provider" @@ -345,6 +346,7 @@ type rewindUndo struct { tail []ai.Message reasoning []provider.MessageReasoning marks []*messagePresentation + instants []time.Time } // RewindKeep uses the terminal's cut and retains one tail until another turn or cut. @@ -372,6 +374,7 @@ func (a *Agent) RewindKeep(index int) (string, error) { undo := &rewindUndo{token: hex.EncodeToString(bytes[:]), tail: append([]ai.Message(nil), a.messages[index:]...), reasoning: append([]provider.MessageReasoning(nil), a.messageReasoning[index:]...)} for _, m := range undo.tail { undo.marks = append(undo.marks, a.presentation.of(m)) + undo.instants = append(undo.instants, a.file.written(m)) } a.cutLocked(index) undo.fingerprint = rewindFingerprint(a.messages) @@ -394,6 +397,7 @@ func (a *Agent) UndoRewind(token string) error { a.presentation.remember(m, u.marks[i]) if a.file != nil { a.file.presentation.remember(m, u.marks[i]) + a.file.rememberStamp(m, u.instants[i]) } } if a.file != nil { diff --git a/internal/session/rewind_test.go b/internal/session/rewind_test.go index 5309025f13..59cbe08231 100644 --- a/internal/session/rewind_test.go +++ b/internal/session/rewind_test.go @@ -5,6 +5,7 @@ import ( "encoding/json" "errors" "path/filepath" + "reflect" "strings" "testing" "time" @@ -700,3 +701,61 @@ func TestRewindKeepUndoSurvivesResume(t *testing.T) { t.Fatal("restored message was not replayed exactly once") } } + +func TestRewindKeepUndoPreservesSteerInstantAndRestampsOrdinaryMessages(t *testing.T) { + const at = "2026-10-15T14:00:00Z" + path := filepath.Join(t.TempDir(), "session.jsonl") + file, _, err := openSessionFile(path, t.TempDir(), "test/model", "") + if err != nil { + t.Fatal(err) + } + for _, line := range []sessionEntry{ + {Type: "message", Role: "user", Content: "start", Timestamp: at}, + {Type: "message", Role: "assistant", Content: "working", ToolCalls: []ai.ToolCall{{ID: "c1", Type: "function", Function: ai.ToolCallFunction{Name: "read", Arguments: "{}"}}}, Timestamp: at}, + {Type: "message", Role: "tool", ToolCallID: "c1", Content: "result", Timestamp: at}, + {Type: "message", Role: "user", Content: "correction", Steer: &journalSteer{}, Timestamp: at}, + {Type: "message", Role: "user", Content: "session note", Note: true, Timestamp: at}, + {Type: "message", Role: "assistant", Content: "done", Timestamp: at}, + } { + file.writeLine(line) + } + if err := file.Close(); err != nil { + t.Fatal(err) + } + a, _ := newTestAgent(t, &refusingCompleter{t}, func(c *Config) { c.SessionFile = path }) + before := a.Transcript() + if len(before) != 7 { + t.Fatalf("resumed %d entries, want 7", len(before)) + } + token, err := a.RewindKeep(a.RewindPoints()[0].Index) + if err != nil { + t.Fatal(err) + } + started := time.Now() + if err := a.UndoRewind(token); err != nil { + t.Fatal(err) + } + finished := time.Now() + after := a.Transcript() + written := ReadTranscript(path).Entries + if len(after) != len(before) || len(written) != len(before) { + t.Fatalf("undo entries: live %d, journal %d, want %d", len(after), len(written), len(before)) + } + for i, entry := range after { + if entry.Text == "correction" { + if !entry.At.Equal(lineInstant(at)) { + t.Fatalf("steer instant = %v, want original journal instant %s", entry.At, at) + } + } else { + // appendWindow has always sent ordinary messages through appendMessage, + // which stamps their new writes rather than preserving the old instant. + if !entry.At.Equal(written[i].At) || entry.At.Before(started) || entry.At.After(finished) { + t.Fatalf("%q instant = %v, want its undo write instant %v within [%v, %v]", entry.Text, entry.At, written[i].At, started, finished) + } + } + after[i].At = before[i].At + } + if !reflect.DeepEqual(after, before) { + t.Fatal("undo changed shaped entries beyond their ordinary write instants") + } +} diff --git a/internal/session/sessionfile.go b/internal/session/sessionfile.go index f8f17688e6..3636aa8af3 100644 --- a/internal/session/sessionfile.go +++ b/internal/session/sessionfile.go @@ -1752,6 +1752,9 @@ func openSessionFile(path, cwd, model, id string) (*sessionFile, replayedSession } journal.restored = replayed.usage journal.remoteSeen = replayed.remoteSeen + if info, err := file.Stat(); err == nil { + replayed.readBytes = info.Size() + } if !replayed.existed { // The header names the session once. A resumed file keeps its @@ -2500,6 +2503,9 @@ func legacyCompactionNote(summary string) string { // the name arrived here after the other two — and a reader of a call site // should not have to count positions to know which bool is which. type replayedSession struct { + // readBytes lets the resume caller release a large replay's working memory + // after shaping, without looking up the journal's path a second time. + readBytes int64 modelSwitches []ModelSwitch // userLines is read-only physical provenance for display reads, nil on resume. userLines map[any]int diff --git a/internal/session/stamp_retain_test.go b/internal/session/stamp_retain_test.go new file mode 100644 index 0000000000..d0f2946dc3 --- /dev/null +++ b/internal/session/stamp_retain_test.go @@ -0,0 +1,339 @@ +package session + +import ( + "context" + "fmt" + "path/filepath" + "reflect" + "runtime" + "strings" + "testing" + "time" + + "github.com/Agent-Field/agentfield/sdk/go/ai" +) + +// Repeated windows use the journal's writer so the fixture follows the same +// marker and re-recording path as a live pass, including occurrence timestamps. +func stampWindowJournal(t *testing.T, passes, width int) string { + t.Helper() + path := filepath.Join(t.TempDir(), "session.jsonl") + file, _, err := openSessionFile(path, t.TempDir(), "test/model", "") + if err != nil { + t.Fatal(err) + } + window := make([]ai.Message, width) + for i := range window { + role := "assistant" + if i%2 == 0 { + role = "user" + } + window[i] = textMessage(role, strings.Repeat("long conversation ", 128)) + file.append(window[i], false, nil) + } + for range passes { + file.appendCompaction(compactionPass{folded: 1}, 1000, window) + } + if err := file.Close(); err != nil { + t.Fatal(err) + } + return path +} + +func stampInstants(a *Agent) map[any]time.Time { + a.file.mu.Lock() + stamps := a.file.stamps + a.file.mu.Unlock() + stamps.mu.Lock() + defer stamps.mu.Unlock() + instants := make(map[any]time.Time, len(stamps.at)) + for key, at := range stamps.at { + instants[key] = at + } + return instants +} + +// Window writes refresh ordinary timestamps before pruning. Read those new +// instants independently; steers and queued messages keep their prior instants. +func journalStampInstants(t *testing.T, a *Agent, prior map[any]time.Time) map[any]time.Time { + t.Helper() + replayed, err := replaySessionFile(a.config.SessionFile) + if err != nil { + t.Fatal(err) + } + a.mu.Lock() + defer a.mu.Unlock() + if len(replayed.messages) != len(a.messages)-1 { + t.Fatalf("journal has %d messages, live window has %d", len(replayed.messages), len(a.messages)-1) + } + want := make(map[any]time.Time, len(prior)+len(replayed.messages)) + for key, at := range prior { + want[key] = at + } + for i, m := range a.messages[1:] { + if chatRefKey(m) != chatRefKey(replayed.messages[i]) { + t.Fatalf("live message %d differs from journal", i) + } + if a.file.steerMark(m) != nil && !want[stampKey(m)].IsZero() { + continue + } + if at := replayed.stamps.of(replayed.messages[i]); !at.IsZero() { + want[stampKey(m)] = at + } + } + return want +} + +func assertLiveStamps(t *testing.T, a *Agent, want map[any]time.Time) { + t.Helper() + a.mu.Lock() + defer a.mu.Unlock() + live := map[any]bool{} + add := func(m ai.Message) { live[stampKey(m)] = true } + for _, m := range a.messages { + add(m) + } + for _, q := range a.steering { + add(q.message) + } + for _, q := range a.ambient { + add(q.message) + } + for _, q := range a.followups { + add(q.message.message) + } + a.file.mu.Lock() + stamps := a.file.stamps + a.file.mu.Unlock() + stamps.mu.Lock() + defer stamps.mu.Unlock() + for key := range stamps.at { + if !live[key] { + t.Fatalf("stamp retains a discarded message: %d stamps for %d live identities", len(stamps.at), len(live)) + } + } + for key := range live { + if at, ok := want[key]; ok && !stamps.at[key].Equal(at) { + t.Fatalf("live stamp = %v, want %v", stamps.at[key], at) + } + } +} + +func TestResumeRetainsOnlyLiveStamps(t *testing.T) { + path := stampWindowJournal(t, 40, 32) + a, _ := newTestAgent(t, &refusingCompleter{t}, func(c *Config) { c.SessionFile = path }) + assertLiveStamps(t, a, journalStampInstants(t, a, nil)) +} + +func TestResumeHeapDoesNotScaleWithDiscardedWindows(t *testing.T) { + measure := func(passes int) uint64 { + path := stampWindowJournal(t, passes, 48) + runtime.GC() + var before, after runtime.MemStats + runtime.ReadMemStats(&before) + a, _ := newTestAgent(t, &refusingCompleter{t}, func(c *Config) { c.SessionFile = path }) + runtime.GC() + runtime.ReadMemStats(&after) + runtime.KeepAlive(a) + if after.HeapAlloc < before.HeapAlloc { + return 0 + } + return after.HeapAlloc - before.HeapAlloc + } + minimum := func(passes int) uint64 { + var samples [3]uint64 + for i := range samples { + samples[i] = measure(passes) + } + best := min(samples[0], samples[1], samples[2]) + t.Logf("%d passes: retained heap samples %v bytes, minimum %d bytes", passes, samples, best) + return best + } + short, long := minimum(6), minimum(60) + const duplicated = 54 * 48 * (len("long conversation ") * 128) + t.Logf("retained heap growth: 6 passes %d bytes, 60 passes %d bytes; duplicated content %d bytes", short, long, duplicated) + if int64(long)-int64(short) >= int64(duplicated/4) { + t.Fatalf("discarded windows retained %d extra bytes; limit %d", int64(long)-int64(short), duplicated/4) + } +} + +func TestLiveCompactionTurnFoldAndRewindRetainOnlyLiveStamps(t *testing.T) { + queueInstants := func(a *Agent) { + a.mu.Lock() + defer a.mu.Unlock() + // Queues are outside the rewritten window, so every live boundary must + // preserve their original instants rather than just their keys. + messages := []ai.Message{textMessage("user", "queued steering"), textMessage("user", "queued ambient"), textMessage("user", "queued followup")} + for _, m := range messages { + a.file.stamps.remember(m, lineInstant("2026-10-15T14:00:00Z")) + } + a.steering = append(a.steering, userMessage{message: messages[0]}) + a.ambient = append(a.ambient, userMessage{message: messages[1]}) + a.followups = append(a.followups, followUp{message: userMessage{message: messages[2]}}) + t.Cleanup(func() { + a.mu.Lock() + a.steering, a.ambient, a.followups = nil, nil, nil + a.mu.Unlock() + }) + } + a, _ := newTestAgent(t, &refusingCompleter{t}, func(c *Config) { + c.SessionFile = filepath.Join(t.TempDir(), "session.jsonl") + c.ContextWindow = 2_000_000 + }) + a.mu.Lock() + for _, m := range exchanges(6, map[int]string{1: strings.Repeat("large tool result\n", 400)}) { + a.messages = append(a.messages, m) + a.file.append(m, false, nil) + } + a.mu.Unlock() + queueInstants(a) + instants := stampInstants(a) + if changed, err := a.compact(context.Background(), nil); !changed || err != nil { + t.Fatalf("compact = %v, %v", changed, err) + } + assertLiveStamps(t, a, journalStampInstants(t, a, instants)) + before := a.Transcript() + if want := ReadTranscript(a.config.SessionFile).Entries; !reflect.DeepEqual(before, want) { + t.Fatal("live compaction entries or times differ from the journal") + } + instants = stampInstants(a) + if _, err := a.Rewind(); err != nil { + t.Fatal(err) + } + assertLiveStamps(t, a, instants) + after := a.Transcript() + if !reflect.DeepEqual(after, before[:len(after)]) { + t.Fatal("rewind changed the retained entries or their times") + } + + folded, _ := newTestAgent(t, &refusingCompleter{t}, func(c *Config) { + c.SessionFile = filepath.Join(t.TempDir(), "session.jsonl") + c.ContextWindow = bigTestWindow + }) + folded.mu.Lock() + folded.turnFloor = len(folded.messages) + for round := range 24 { + appendToolRound(folded, fmt.Sprintf("call-%d", round), "read", "{}", turnFoldOutput()) + } + for _, m := range folded.messages[1:] { + folded.file.append(m, false, nil) + } + folded.mu.Unlock() + queueInstants(folded) + instants = stampInstants(folded) + foldWholeTurn(t, folded) + if !holdsPrefix(folded.snapshot(), foldMarkerPrefix) { + t.Fatal("turn-fold fixture did not reach the pruning boundary") + } + assertLiveStamps(t, folded, journalStampInstants(t, folded, instants)) + if want := ReadTranscript(folded.config.SessionFile).Entries; !reflect.DeepEqual(folded.Transcript(), want) { + t.Fatal("live turn-fold entries or times differ from the journal") + } +} + +func TestResumeStampPrunePreservesShapedEntries(t *testing.T) { + const at = "2026-10-15T14:00:00Z" + message := func(role, text string) sessionEntry { + return sessionEntry{Type: "message", Role: role, Content: text, Timestamp: at} + } + tools := []sessionEntry{message("user", "inspect"), { + Type: "message", Role: "assistant", ToolCalls: []ai.ToolCall{{ID: "c1", Type: "function", Function: ai.ToolCallFunction{Name: "read", Arguments: "{}"}}}, Timestamp: at, + }, {Type: "message", Role: "tool", ToolCallID: "c1", Content: "result", Timestamp: at}, message("assistant", "done")} + for name, lines := range map[string][]sessionEntry{ + "empty": nil, + "short": {message("user", "hello"), message("assistant", "hi")}, + "tools": tools, + "image and attachment": {{Type: "message", Role: "user", Content: "see image", Timestamp: at, + Parts: []journalPart{{Type: journalPartImage, Path: "/tmp/stamp-test.png", MIME: "image/png", SHA256: strings.Repeat("a", 64)}}, + Presentation: &messagePresentation{Audience: audiencePerson, Attachments: []MessageAttachment{{Name: "stamp-test.png", Path: "/tmp/stamp-test.png"}}}}}, + "rewind": {message("user", "keep"), message("assistant", "kept"), message("user", "drop"), {Type: "rewind", Dropped: 1}, message("user", "replacement")}, + "steer": {message("user", "start"), message("assistant", "working"), {Type: "message", Role: "user", Content: "correction", Timestamp: at, Steer: &journalSteer{}}}, + "compactions": {message("user", "first"), message("assistant", "old"), {Type: "compaction", Folded: 1, Window: 2, Timestamp: at}, message("user", "first"), message("assistant", "folded"), {Type: "compaction", Folded: 1, Window: 2, Timestamp: at}, message("user", "first"), message("assistant", "folded again")}, + } { + t.Run(name, func(t *testing.T) { + path := filepath.Join(t.TempDir(), "session.jsonl") + file, _, err := openSessionFile(path, t.TempDir(), "test/model", "") + if err != nil { + t.Fatal(err) + } + for _, line := range lines { + file.writeLine(line) + } + if err := file.Close(); err != nil { + t.Fatal(err) + } + // Shape the resume replay before the constructor can prune its indexes. + file, replayed, err := openSessionFile(path, t.TempDir(), "test/model", "") + if err != nil { + t.Fatal(err) + } + want := shapeEntries(replayed.messages, file) + earlier := shapeEntries(replayed.earlier, file) + floor := len(shapeEntries(replayed.messages[:replayed.overlap], file)) + file.Close() + a, _ := newTestAgent(t, &refusingCompleter{t}, func(c *Config) { + c.SessionFile = path + c.SupportsImages = func(string) bool { return true } + }) + if !reflect.DeepEqual(a.Transcript(), want) || !reflect.DeepEqual(a.EarlierHistory(), EarlierHistory{Entries: earlier, Floor: floor}) { + t.Fatal("pruning changed shaped entries, identities, metadata or times") + } + for _, entry := range a.Transcript() { + if entry.At.IsZero() || !entry.At.Equal(lineInstant(at)) { + t.Fatalf("entry lost its journal instant: %+v", entry) + } + } + }) + } +} + +func TestStampRetainKeepsQueueTimes(t *testing.T) { + a, _ := newTestAgent(t, &refusingCompleter{t}, func(c *Config) { c.SessionFile = filepath.Join(t.TempDir(), "session.jsonl") }) + // A queued message can already be journaled; its stamp must survive a live + // boundary before the message joins the model's transcript. + a.mu.Lock() + defer a.mu.Unlock() + messages := []ai.Message{textMessage("user", "steering"), textMessage("user", "ambient"), textMessage("user", "followup"), textMessage("assistant", "retained reply")} + ats := make([]time.Time, len(messages)) + for i, m := range messages { + a.file.append(m, false, nil) + ats[i] = a.file.written(m) + } + a.messages = append(a.messages, messages[3]) + a.steering = append(a.steering, userMessage{message: messages[0]}) + a.ambient = append(a.ambient, userMessage{message: messages[1]}) + a.followups = append(a.followups, followUp{message: userMessage{message: messages[2]}}) + defer func() { + a.steering, a.ambient, a.followups = nil, nil, nil + }() + a.retainPresentationLocked() + for i, m := range messages { + if got := a.file.written(m); got.IsZero() || !got.Equal(ats[i]) { + t.Fatalf("queued instant = %v, want %v", got, ats[i]) + } + } +} + +func TestStampRetainKeepsOnlyRequestedInstants(t *testing.T) { + var absent *stampIndex + absent.retain(nil) + var noFile *sessionFile + noFile.retainStamps(nil) + s := &stampIndex{} + s.retain(nil) + at := time.Date(2026, 10, 15, 14, 0, 0, 0, time.UTC) + word := textMessage("user", "keep") + call := ai.Message{Role: "assistant", ToolCalls: []ai.ToolCall{{ID: "call-only"}}} + dropped := textMessage("assistant", "gone") + for _, m := range []ai.Message{word, call, dropped} { + s.remember(m, at) + } + s.retain([]ai.Message{word}, []ai.Message{call, {}}) + if len(s.at) != 2 || !s.of(word).Equal(at) || !s.of(call).Equal(at) || !s.of(dropped).IsZero() { + t.Fatal("retain lost a live instant or kept a discarded one") + } + s.retain() + if len(s.at) != 0 { + t.Fatal("retain with no messages kept stamps") + } +}