session: an opened chat lets go of the journal lines it replayed

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.

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

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 {

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()
+	}()
+}