From 49ff4c0678683b942e2c463ebde5a6475595ed98 Mon Sep 17 00:00:00 2001
From: Ed Zynda <ezynda3@gmail.com>
Date: Sun, 29 Mar 2026 15:02:24 +0300
Subject: [PATCH] fix: /tree and /fork commands lose context due to leaf reset
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

performFork() called ClearMessages() after Branch(targetID), but
ClearMessages() calls TreeSession.ResetLeaf() which sets leafID back
to empty — immediately undoing the branch. The in-memory message store
was also never reloaded from the tree session after branching, so the
LLM had zero context.

Add ReloadMessagesFromTree() which clears the store and reloads it
from the tree session's current branch without resetting the leaf
pointer. Use it in performFork() instead of ClearMessages().
---
 internal/app/app.go       | 11 +++++++++++
 internal/ui/model.go      | 11 ++++++++++-
 internal/ui/model_test.go |  4 ++++
 3 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/internal/app/app.go b/internal/app/app.go
index 80cac5ec..44579ece 100644
--- a/internal/app/app.go
+++ b/internal/app/app.go
@@ -271,6 +271,17 @@ func (a *App) ClearMessages() {
 	}
 }
 
+// ReloadMessagesFromTree clears the in-memory message store and reloads it
+// from the tree session's current branch. Unlike ClearMessages, this does NOT
+// reset the tree session's leaf pointer. Used after Branch() to sync the
+// store with the new branch position.
+func (a *App) ReloadMessagesFromTree() {
+	a.store.Clear()
+	if a.opts.TreeSession != nil {
+		a.store.Replace(a.opts.TreeSession.GetLLMMessages())
+	}
+}
+
 // GetTreeSession returns the tree session manager, or nil if not configured.
 func (a *App) GetTreeSession() *session.TreeManager {
 	return a.opts.TreeSession
diff --git a/internal/ui/model.go b/internal/ui/model.go
index d22c5796..6d4f00a7 100644
--- a/internal/ui/model.go
+++ b/internal/ui/model.go
@@ -74,6 +74,11 @@ type AppController interface {
 	ClearQueue()
 	// ClearMessages clears the conversation history.
 	ClearMessages()
+	// ReloadMessagesFromTree clears the in-memory message store and reloads
+	// it from the tree session's current branch. Unlike ClearMessages, this
+	// does NOT reset the tree session's leaf pointer. Used after Branch() to
+	// sync the store with the new branch position.
+	ReloadMessagesFromTree()
 	// CompactConversation summarises older messages to free context space.
 	// Runs asynchronously; results are delivered via CompactCompleteEvent or
 	// CompactErrorEvent sent through the registered tea.Program. Returns an
@@ -3047,8 +3052,12 @@ func (m *AppModel) performFork(targetID string, isUser bool, userText string) te
 		return nil
 	}
 
+	// Branch the tree session to the target entry. We must NOT call
+	// ClearMessages() here because it resets the leaf pointer back to "",
+	// undoing the branch we just set. Instead, branch first and then
+	// reload the in-memory store from the tree session's current branch.
 	_ = ts.Branch(targetID)
-	m.appCtrl.ClearMessages()
+	m.appCtrl.ReloadMessagesFromTree()
 
 	// If it was a user message, populate the input with the text.
 	if isUser && userText != "" {
diff --git a/internal/ui/model_test.go b/internal/ui/model_test.go
index ee0a69b8..62e0eb9e 100644
--- a/internal/ui/model_test.go
+++ b/internal/ui/model_test.go
@@ -46,6 +46,10 @@ func (s *stubAppController) ClearMessages() {
 	s.clearMsgCalled++
 }
 
+func (s *stubAppController) ReloadMessagesFromTree() {
+	// no-op in tests
+}
+
 func (s *stubAppController) CompactConversation(_ string) error {
 	return nil
 }