From 822b318b18b46e70351accc5508b31678f05a685 Mon Sep 17 00:00:00 2001 From: pdfinn Date: Thu, 26 Feb 2026 06:00:53 +0700 Subject: [PATCH] fix(llmfs): clean up sessions on disconnect to prevent memory leak MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Each spawned subagent creates an LLM session via /n/llm/new, but sessions were never automatically freed: BaseFile.Close() was a no-op, and connection drop handling only removed the client from the map without touching sessions. Repeated spawn calls caused sessions to accumulate indefinitely, consuming memory for full conversation histories. Changes: - Add refs int32 (atomic) to Session struct - Add IncRef(id) / DecRef(id) to SessionManager; DecRef calls Close when refs reach zero, freeing the session and all its conversation history - Add sessionRefFile / sessionRefDir wrappers (session_ref.go) that call IncRef on Open and DecRef on Close (once.Do guards against double-decrement) - SessionsDir.Lookup wraps the returned SessionDir in sessionRefDir, so child file lookups via sessionRefDir.Lookup also yield sessionRefFile objects - Fix handleConn disconnect path to call file.Close() on all unclosed fids before removing the client, ensuring sessions are freed on network drop Typical flow: subagent exits → NEWFD closes all FDs → 9P connection drops → handleConn deferred cleanup calls Close on remaining fids → sessionRefFile.Close → sm.DecRef → refs=0 → sm.Close(id) → session deleted from map. Co-Authored-By: Claude Sonnet 4.6 --- internal/llm/session.go | 25 ++++++++++++ internal/llmfs/session_dir.go | 9 +++- internal/llmfs/session_ref.go | 77 +++++++++++++++++++++++++++++++++++ internal/protocol/server.go | 7 ++++ 4 files changed, 117 insertions(+), 1 deletion(-) create mode 100644 internal/llmfs/session_ref.go diff --git a/internal/llm/session.go b/internal/llm/session.go index b46ba6561076d67bd32e4d5b24d90cd9c639a14d..f0459962fef446b890bbaadd00591b2df65ba3e4 100644 --- a/internal/llm/session.go +++ b/internal/llm/session.go @@ -7,6 +7,7 @@ import ( "fmt" "strings" "sync" + "sync/atomic" ) // SessionDefaults are copied to new sessions at creation time. @@ -48,6 +49,7 @@ type Session struct { mu sync.RWMutex closed bool + refs int32 // atomic reference count; session closed when it drops to 0 // Async generation + streaming support streamCh chan string // raw text chunks during generation; nil when idle @@ -377,6 +379,29 @@ func (sm *SessionManager) Close(id int) error { return nil } +// IncRef increments the reference count for the session with the given ID. +// Called when a 9P client opens a file inside the session directory. +func (sm *SessionManager) IncRef(id int) { + session := sm.Get(id) + if session != nil { + atomic.AddInt32(&session.refs, 1) + } +} + +// DecRef decrements the reference count for the session with the given ID. +// When the count reaches zero the session is closed and removed, freeing +// all conversation history from memory. Called when a 9P client clunks +// (closes) a file inside the session directory. +func (sm *SessionManager) DecRef(id int) { + session := sm.Get(id) + if session == nil { + return + } + if atomic.AddInt32(&session.refs, -1) <= 0 { + sm.Close(id) //nolint:errcheck + } +} + // Reset clears the conversation history for the given session. func (sm *SessionManager) Reset(id int) error { session := sm.Get(id) diff --git a/internal/llmfs/session_dir.go b/internal/llmfs/session_dir.go index 6eb9c6c0efd6f4e0d0df01f9ee87d6ffc959c556..3feaa2daf052a867aa75f3d451e4d3840ba7e330 100644 --- a/internal/llmfs/session_dir.go +++ b/internal/llmfs/session_dir.go @@ -160,7 +160,14 @@ func (d *SessionsDir) Lookup(name string) (protocol.File, error) { return nil, protocol.ErrNotFound } - return NewSessionDir(d.sm, id), nil + // Wrap in sessionRefDir so that any Topen of the directory (or its + // children via sessionRefDir.Lookup) increments the session ref count. + // When the last ref is clunked the session is freed automatically. + return &sessionRefDir{ + Dir: NewSessionDir(d.sm, id), + sm: d.sm, + id: id, + }, nil } // Read returns directory listing as packed stat entries. diff --git a/internal/llmfs/session_ref.go b/internal/llmfs/session_ref.go new file mode 100644 index 0000000000000000000000000000000000000000..4879ea30b94c50d7b11bb22c4917ad88823f3e1f --- /dev/null +++ b/internal/llmfs/session_ref.go @@ -0,0 +1,77 @@ +package llmfs + +import ( + "sync" + + "github.com/NERVsystems/llm9p/internal/llm" + "github.com/NERVsystems/llm9p/internal/protocol" +) + +// sessionRefFile wraps a protocol.File and tracks a reference to its session. +// It calls sm.IncRef(id) on Open and sm.DecRef(id) on the first Close, so +// that sessions are automatically cleaned up when all their files are clunked. +type sessionRefFile struct { + protocol.File + sm *llm.SessionManager + id int + once sync.Once + opened bool +} + +func (r *sessionRefFile) Open(mode uint8) error { + err := r.File.Open(mode) + if err == nil { + r.opened = true + r.sm.IncRef(r.id) + } + return err +} + +func (r *sessionRefFile) Close() error { + err := r.File.Close() + if r.opened { + r.once.Do(func() { + r.sm.DecRef(r.id) + }) + } + return err +} + +// sessionRefDir wraps a protocol.Dir (SessionDir) with the same ref counting. +// It additionally overrides Lookup so that child files are also ref-counted. +type sessionRefDir struct { + protocol.Dir + sm *llm.SessionManager + id int + once sync.Once + opened bool +} + +func (r *sessionRefDir) Open(mode uint8) error { + err := r.Dir.Open(mode) + if err == nil { + r.opened = true + r.sm.IncRef(r.id) + } + return err +} + +func (r *sessionRefDir) Close() error { + err := r.Dir.Close() + if r.opened { + r.once.Do(func() { + r.sm.DecRef(r.id) + }) + } + return err +} + +// Lookup wraps child files in sessionRefFile so each open fid contributes +// its own reference to the session. +func (r *sessionRefDir) Lookup(name string) (protocol.File, error) { + f, err := r.Dir.Lookup(name) + if err != nil { + return nil, err + } + return &sessionRefFile{File: f, sm: r.sm, id: r.id}, nil +} diff --git a/internal/protocol/server.go b/internal/protocol/server.go index 6eb5bab9ddc18879a2bf9871dff48828b72f0415..dd7074a4aeeb6bf3684eefc9c942d9e9b9b32756 100644 --- a/internal/protocol/server.go +++ b/internal/protocol/server.go @@ -72,6 +72,13 @@ func (s *Server) handleConn(conn net.Conn) { s.mu.Unlock() defer func() { + // Close all fids that were not explicitly clunked before disconnect. + // This triggers sessionRefFile/sessionRefDir.Close() which decrements + // session reference counts, allowing sessions to be freed when the + // last client using them disconnects. + for _, file := range state.fids { + file.Close() //nolint:errcheck + } s.mu.Lock() delete(s.clients, conn) s.mu.Unlock()