From bbf6fc70a628186f760ea688c6ff7feda9a479ac Mon Sep 17 00:00:00 2001
From: Ed Zynda <ezynda3@gmail.com>
Date: Thu, 30 Jul 2026 15:33:44 +0300
Subject: [PATCH 1/2] refactor(app): expose session state as value snapshots
 (#101)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The TUI reached through AppController.GetTreeSession() to grab the live
*session.TreeManager and drove it directly: reading the header, walking
the branch, appending entries, forking, and type-switching on the
concrete *session.MessageEntry / *session.ModelChangeEntry / ... entry
types to decide how to filter and render the tree. That made the session
storage schema part of the UI's contract — any change to how entries are
persisted broke rendering, and the UI could not be exercised or reused
without a real session on disk.

Session state is now projected into immutable value types owned by the
app layer:

  - SessionSnapshot carries session metadata (id, name, file path, cwd,
    created, leaf, counts, persisted).
  - TreeNodeView carries a node's id, parent, kind, role, text payload,
    label and children. EntryKind replaces the concrete-type switches,
    so an unrecognised entry degrades to a fallback instead of breaking.
  - SessionHistory returns the current branch as []message.Message, so
    transcript replay no longer decodes entries itself.

Mutations move behind intent-revealing methods (SetSessionName,
NewSession, ForkSession, SessionSystemPromptEntry) instead of the UI
constructing sessions and swapping them in via SwitchTreeSession.

GetTreeSession and SwitchTreeSession are dropped from AppController;
App keeps both for cmd/root.go wiring, which owns session lifecycle.

TreeNodeSelectedMsg now carries ParentID directly rather than the raw
entry, removing the entry lookup and type assertion the fork path did to
find where to branch from.

Behaviour is unchanged: filters, display strings, colours, export/share
filenames and fork semantics all match the previous implementation, now
covered by tests for the projections and for tree filtering/rendering.

The session selector still calls session.ListSessions/DeleteSession
directly; moving session discovery and the export/share file I/O out of
the UI is the next step.

Refs #101
---
 internal/app/session_view.go      | 284 +++++++++++++++++++++++++
 internal/app/session_view_test.go | 337 ++++++++++++++++++++++++++++++
 internal/session/entry.go         |  27 +++
 internal/ui/core/events.go        |   6 +-
 internal/ui/model.go              | 217 +++++++++----------
 internal/ui/model_test.go         |  28 ++-
 internal/ui/tree_selector.go      | 156 ++++++--------
 internal/ui/tree_selector_test.go | 231 ++++++++++++++++++++
 8 files changed, 1072 insertions(+), 214 deletions(-)
 create mode 100644 internal/app/session_view.go
 create mode 100644 internal/app/session_view_test.go
 create mode 100644 internal/ui/tree_selector_test.go

diff --git a/internal/app/session_view.go b/internal/app/session_view.go
new file mode 100644
index 00000000..b340dc74
--- /dev/null
+++ b/internal/app/session_view.go
@@ -0,0 +1,284 @@
+package app
+
+import (
+	"errors"
+	"time"
+
+	"github.com/mark3labs/kit/internal/message"
+	"github.com/mark3labs/kit/internal/session"
+)
+
+// ErrNoSession is returned by session mutation methods when no tree session
+// is active. Callers that only need to know whether a session exists should
+// prefer the ok result of SessionSnapshot.
+var ErrNoSession = errors.New("no tree session active")
+
+// EntryKind classifies a session tree entry without exposing the concrete
+// persistence types from internal/session.
+//
+// Presentation layers switch on EntryKind instead of type-switching on
+// *session.MessageEntry and friends, so the on-disk entry schema can evolve
+// without breaking every consumer. Unrecognised entries carry
+// EntryKindUnknown; consumers should always handle it.
+type EntryKind string
+
+const (
+	// EntryKindUnknown is used for entries the app layer does not recognise.
+	EntryKindUnknown EntryKind = ""
+	// EntryKindMessage is a conversation message (user, assistant or tool).
+	EntryKindMessage EntryKind = "message"
+	// EntryKindModelChange records a provider/model switch.
+	EntryKindModelChange EntryKind = "model_change"
+	// EntryKindBranchSummary carries a summary of an abandoned branch.
+	EntryKindBranchSummary EntryKind = "branch_summary"
+	// EntryKindLabel bookmarks another entry with a user-defined label.
+	EntryKindLabel EntryKind = "label"
+	// EntryKindSessionInfo records the session's display name.
+	EntryKindSessionInfo EntryKind = "session_info"
+	// EntryKindExtensionData holds extension-defined persisted state.
+	EntryKindExtensionData EntryKind = "extension_data"
+	// EntryKindCompaction records an LLM-generated summary of older messages.
+	EntryKindCompaction EntryKind = "compaction"
+	// EntryKindSystemPrompt captures the system prompt used for a session.
+	EntryKindSystemPrompt EntryKind = "system_prompt"
+)
+
+// TreeNodeView is an immutable value projection of a single node in the
+// session entry tree. It carries everything a presentation layer needs to
+// filter, search and render the node without reaching back into the session
+// store or knowing the entry schema.
+type TreeNodeView struct {
+	// ID is the entry ID.
+	ID string
+	// ParentID is the parent entry ID, empty for root entries.
+	ParentID string
+	// Kind classifies the entry.
+	Kind EntryKind
+	// Role is the message role ("user", "assistant", "tool") for
+	// EntryKindMessage nodes and empty for every other kind.
+	Role string
+	// Text is the kind-appropriate plain-text payload: message text for
+	// messages, "provider/model" for model changes, the summary for branch
+	// summaries and compactions, the label text for labels and the session
+	// name for session info. Empty for kinds with no textual payload.
+	//
+	// Text is never truncated or collapsed — formatting is the caller's job.
+	Text string
+	// Label is the user-defined label bookmarking this entry, or empty.
+	Label string
+	// Children are the node's child entries, in insertion order.
+	Children []TreeNodeView
+}
+
+// IsUserMessage reports whether the node is a message authored by the user.
+func (n TreeNodeView) IsUserMessage() bool {
+	return n.Kind == EntryKindMessage && n.Role == "user"
+}
+
+// SessionSnapshot is a point-in-time value view of the active tree session's
+// metadata. It is a copy: mutating it has no effect on the session, and it
+// does not go stale in a way that can corrupt state — callers re-read it
+// whenever they need current values.
+type SessionSnapshot struct {
+	// ID is the session UUID.
+	ID string
+	// Name is the user-defined display name, empty if unnamed.
+	Name string
+	// FilePath is the JSONL file backing the session, empty when in-memory.
+	FilePath string
+	// Cwd is the working directory the session was created in.
+	Cwd string
+	// Created is the session creation timestamp.
+	Created time.Time
+	// LeafID is the current position in the entry tree.
+	LeafID string
+	// EntryCount is the total number of entries in the session.
+	EntryCount int
+	// MessageCount is the number of message entries in the session.
+	MessageCount int
+	// Persisted reports whether the session is backed by a file on disk.
+	Persisted bool
+}
+
+// SessionSnapshot returns a value snapshot of the active tree session's
+// metadata. ok is false when no tree session is configured, in which case the
+// returned snapshot is the zero value.
+func (a *App) SessionSnapshot() (SessionSnapshot, bool) {
+	tm := a.opts.TreeSession
+	if tm == nil {
+		return SessionSnapshot{}, false
+	}
+	header := tm.GetHeader()
+	return SessionSnapshot{
+		ID:           header.ID,
+		Name:         tm.GetSessionName(),
+		FilePath:     tm.GetFilePath(),
+		Cwd:          header.Cwd,
+		Created:      header.Timestamp,
+		LeafID:       tm.GetLeafID(),
+		EntryCount:   tm.EntryCount(),
+		MessageCount: tm.MessageCount(),
+		Persisted:    tm.IsPersisted(),
+	}, true
+}
+
+// SessionTree returns the session's entry tree projected into value nodes,
+// rooted at the entries with no parent. Returns nil when no tree session is
+// active.
+//
+// The projection is a deep copy taken under the session's read lock, so the
+// result is safe to hold and walk while the session keeps mutating.
+func (a *App) SessionTree() []TreeNodeView {
+	tm := a.opts.TreeSession
+	if tm == nil {
+		return nil
+	}
+	roots := tm.GetTree()
+	if len(roots) == 0 {
+		return nil
+	}
+	out := make([]TreeNodeView, 0, len(roots))
+	for _, root := range roots {
+		out = append(out, treeNodeView(tm, root))
+	}
+	return out
+}
+
+// treeNodeView recursively projects a session tree node into a TreeNodeView.
+func treeNodeView(tm *session.TreeManager, node *session.TreeNode) TreeNodeView {
+	view := TreeNodeView{
+		ID:       node.ID,
+		ParentID: node.ParentID,
+		Label:    tm.GetLabel(node.ID),
+	}
+	view.Kind, view.Role, view.Text = describeEntry(node.Entry)
+
+	if len(node.Children) > 0 {
+		view.Children = make([]TreeNodeView, 0, len(node.Children))
+		for _, child := range node.Children {
+			view.Children = append(view.Children, treeNodeView(tm, child))
+		}
+	}
+	return view
+}
+
+// describeEntry maps a concrete session entry to its kind, role and textual
+// payload. This is the single place in the codebase that needs to know the
+// full set of entry types for presentation purposes.
+func describeEntry(entry any) (kind EntryKind, role, text string) {
+	switch e := entry.(type) {
+	case *session.MessageEntry:
+		return EntryKindMessage, e.Role, e.Text()
+	case *session.ModelChangeEntry:
+		return EntryKindModelChange, "", e.Provider + "/" + e.ModelID
+	case *session.BranchSummaryEntry:
+		return EntryKindBranchSummary, "", e.Summary
+	case *session.CompactionEntry:
+		return EntryKindCompaction, "", e.Summary
+	case *session.LabelEntry:
+		return EntryKindLabel, "", e.Label
+	case *session.SessionInfoEntry:
+		return EntryKindSessionInfo, "", e.Name
+	case *session.ExtensionDataEntry:
+		return EntryKindExtensionData, "", ""
+	case *session.SystemPromptEntry:
+		return EntryKindSystemPrompt, "", e.Content
+	default:
+		return EntryKindUnknown, "", ""
+	}
+}
+
+// SessionHistory returns the conversation messages on the session's current
+// branch, oldest first. Entries that are not messages, and messages that fail
+// to decode, are skipped. Returns nil when no tree session is active.
+//
+// This is what the UI replays to rebuild the transcript after resuming or
+// forking a session.
+func (a *App) SessionHistory() []message.Message {
+	tm := a.opts.TreeSession
+	if tm == nil {
+		return nil
+	}
+	branch := tm.GetBranch("")
+	if len(branch) == 0 {
+		return nil
+	}
+	out := make([]message.Message, 0, len(branch))
+	for _, entry := range branch {
+		me, ok := entry.(*session.MessageEntry)
+		if !ok {
+			continue
+		}
+		msg, err := me.ToMessage()
+		if err != nil {
+			continue
+		}
+		out = append(out, msg)
+	}
+	return out
+}
+
+// SetSessionName sets the session's display name, persisting it as a
+// session_info entry. Returns ErrNoSession when no tree session is active.
+func (a *App) SetSessionName(name string) error {
+	tm := a.opts.TreeSession
+	if tm == nil {
+		return ErrNoSession
+	}
+	_, err := tm.AppendSessionInfo(name)
+	return err
+}
+
+// NewSession replaces the active tree session with a brand new one created in
+// cwd, closing and flushing the old session. The in-memory message store is
+// reset to the (empty) new session.
+//
+// Returns ErrNoSession when no tree session is active — callers that want to
+// support the session-less mode should check SessionSnapshot first and clear
+// their own state instead.
+func (a *App) NewSession(cwd string) error {
+	if a.opts.TreeSession == nil {
+		return ErrNoSession
+	}
+	ts, err := session.CreateTreeSession(cwd)
+	if err != nil {
+		return err
+	}
+	a.SwitchTreeSession(ts)
+	return nil
+}
+
+// ForkSession creates a new session in cwd containing the history up to and
+// including targetID, then switches to it. The original session is left
+// untouched on disk. Returns ErrNoSession when no tree session is active.
+func (a *App) ForkSession(cwd, targetID string) error {
+	tm := a.opts.TreeSession
+	if tm == nil {
+		return ErrNoSession
+	}
+	ts, err := tm.ForkToNewSession(cwd, targetID)
+	if err != nil {
+		return err
+	}
+	a.SwitchTreeSession(ts)
+	return nil
+}
+
+// SessionSystemPromptEntry marshals a system-prompt entry describing the
+// given system prompt together with the model and provider currently in
+// effect for the session. It is embedded in exported and shared session files
+// so a reader can reconstruct the context the conversation ran under.
+//
+// fallbackModelID is used when the session records no model change of its
+// own. Returns ErrNoSession when no tree session is active.
+func (a *App) SessionSystemPromptEntry(systemPrompt, fallbackModelID string) ([]byte, error) {
+	tm := a.opts.TreeSession
+	if tm == nil {
+		return nil, ErrNoSession
+	}
+	_, provider, modelID := tm.BuildContext()
+	if modelID == "" {
+		modelID = fallbackModelID
+	}
+	return session.MarshalEntry(session.NewSystemPromptEntry(systemPrompt, modelID, provider))
+}
diff --git a/internal/app/session_view_test.go b/internal/app/session_view_test.go
new file mode 100644
index 00000000..ac5ebe79
--- /dev/null
+++ b/internal/app/session_view_test.go
@@ -0,0 +1,337 @@
+package app
+
+import (
+	"errors"
+	"testing"
+
+	"github.com/mark3labs/kit/internal/message"
+	"github.com/mark3labs/kit/internal/session"
+)
+
+// newSessionApp creates an App backed by an in-memory tree session.
+func newSessionApp(t *testing.T) (*App, *session.TreeManager) {
+	t.Helper()
+	tm := session.InMemoryTreeSession(t.TempDir())
+	a := New(Options{TreeSession: tm}, nil)
+	return a, tm
+}
+
+// appendMessage appends a single-text-part message with the given role.
+func appendMessage(t *testing.T, tm *session.TreeManager, role message.MessageRole, text string) string {
+	t.Helper()
+	id, err := tm.AppendMessage(message.Message{
+		Role:  role,
+		Parts: []message.ContentPart{message.TextContent{Text: text}},
+	})
+	if err != nil {
+		t.Fatalf("AppendMessage: %v", err)
+	}
+	return id
+}
+
+func appendUserMessage(t *testing.T, tm *session.TreeManager, text string) string {
+	t.Helper()
+	return appendMessage(t, tm, message.RoleUser, text)
+}
+
+// --------------------------------------------------------------------------
+// SessionSnapshot
+// --------------------------------------------------------------------------
+
+func TestSessionSnapshotNoSession(t *testing.T) {
+	a := New(Options{}, nil)
+
+	snap, ok := a.SessionSnapshot()
+	if ok {
+		t.Fatal("expected ok=false when no tree session is configured")
+	}
+	if snap != (SessionSnapshot{}) {
+		t.Fatalf("expected zero snapshot, got %+v", snap)
+	}
+}
+
+func TestSessionSnapshotReflectsSession(t *testing.T) {
+	a, tm := newSessionApp(t)
+
+	appendUserMessage(t, tm, "hello")
+	if _, err := tm.AppendSessionInfo("my session"); err != nil {
+		t.Fatalf("AppendSessionInfo: %v", err)
+	}
+
+	snap, ok := a.SessionSnapshot()
+	if !ok {
+		t.Fatal("expected ok=true with an active tree session")
+	}
+	if snap.ID != tm.GetSessionID() {
+		t.Errorf("ID = %q, want %q", snap.ID, tm.GetSessionID())
+	}
+	if snap.Name != "my session" {
+		t.Errorf("Name = %q, want %q", snap.Name, "my session")
+	}
+	if snap.MessageCount != 1 {
+		t.Errorf("MessageCount = %d, want 1", snap.MessageCount)
+	}
+	if snap.EntryCount != 2 {
+		t.Errorf("EntryCount = %d, want 2", snap.EntryCount)
+	}
+	if snap.LeafID != tm.GetLeafID() {
+		t.Errorf("LeafID = %q, want %q", snap.LeafID, tm.GetLeafID())
+	}
+	if snap.Persisted {
+		t.Error("Persisted = true, want false for an in-memory session")
+	}
+	if snap.FilePath != "" {
+		t.Errorf("FilePath = %q, want empty for an in-memory session", snap.FilePath)
+	}
+}
+
+// A snapshot is a value: mutating the session afterwards must not change the
+// copy the caller already holds.
+func TestSessionSnapshotIsDetached(t *testing.T) {
+	a, tm := newSessionApp(t)
+	appendUserMessage(t, tm, "first")
+
+	before, _ := a.SessionSnapshot()
+	appendUserMessage(t, tm, "second")
+
+	if before.MessageCount != 1 {
+		t.Errorf("held snapshot mutated: MessageCount = %d, want 1", before.MessageCount)
+	}
+	after, _ := a.SessionSnapshot()
+	if after.MessageCount != 2 {
+		t.Errorf("fresh snapshot stale: MessageCount = %d, want 2", after.MessageCount)
+	}
+}
+
+// --------------------------------------------------------------------------
+// SessionTree
+// --------------------------------------------------------------------------
+
+func TestSessionTreeNoSession(t *testing.T) {
+	a := New(Options{}, nil)
+	if tree := a.SessionTree(); tree != nil {
+		t.Fatalf("expected nil tree without a session, got %+v", tree)
+	}
+}
+
+func TestSessionTreeProjectsEntryKinds(t *testing.T) {
+	a, tm := newSessionApp(t)
+
+	userID := appendUserMessage(t, tm, "what is 2+2")
+	appendMessage(t, tm, message.RoleAssistant, "4")
+	if _, err := tm.AppendModelChange("anthropic", "claude-sonnet-4-5"); err != nil {
+		t.Fatalf("AppendModelChange: %v", err)
+	}
+	if _, err := tm.AppendLabel(userID, "the question"); err != nil {
+		t.Fatalf("AppendLabel: %v", err)
+	}
+
+	// Flatten the projected tree so we can assert on it regardless of shape.
+	var flat []TreeNodeView
+	var walk func([]TreeNodeView)
+	walk = func(nodes []TreeNodeView) {
+		for _, n := range nodes {
+			flat = append(flat, n)
+			walk(n.Children)
+		}
+	}
+	walk(a.SessionTree())
+
+	byID := make(map[string]TreeNodeView, len(flat))
+	for _, n := range flat {
+		byID[n.ID] = n
+	}
+
+	user, ok := byID[userID]
+	if !ok {
+		t.Fatalf("user message %q missing from projected tree", userID)
+	}
+	if user.Kind != EntryKindMessage {
+		t.Errorf("user Kind = %q, want %q", user.Kind, EntryKindMessage)
+	}
+	if user.Role != "user" {
+		t.Errorf("user Role = %q, want %q", user.Role, "user")
+	}
+	if user.Text != "what is 2+2" {
+		t.Errorf("user Text = %q, want %q", user.Text, "what is 2+2")
+	}
+	if !user.IsUserMessage() {
+		t.Error("IsUserMessage() = false for a user message")
+	}
+	// Labels applied to an entry are surfaced on the node they target.
+	if user.Label != "the question" {
+		t.Errorf("user Label = %q, want %q", user.Label, "the question")
+	}
+
+	var sawModelChange bool
+	for _, n := range flat {
+		if n.Kind == EntryKindModelChange {
+			sawModelChange = true
+			if n.Text != "anthropic/claude-sonnet-4-5" {
+				t.Errorf("model change Text = %q, want %q", n.Text, "anthropic/claude-sonnet-4-5")
+			}
+			if n.Role != "" {
+				t.Errorf("model change Role = %q, want empty", n.Role)
+			}
+			if n.IsUserMessage() {
+				t.Error("IsUserMessage() = true for a model change entry")
+			}
+		}
+	}
+	if !sawModelChange {
+		t.Error("model change entry missing from projected tree")
+	}
+}
+
+// The projection must be a deep copy: mutating it must not be observable, and
+// later session writes must not be visible in an already-taken projection.
+func TestSessionTreeIsSnapshot(t *testing.T) {
+	a, tm := newSessionApp(t)
+	appendUserMessage(t, tm, "first")
+
+	before := a.SessionTree()
+	if len(before) != 1 {
+		t.Fatalf("expected 1 root, got %d", len(before))
+	}
+	if len(before[0].Children) != 0 {
+		t.Fatalf("expected no children, got %d", len(before[0].Children))
+	}
+
+	appendUserMessage(t, tm, "second")
+
+	if len(before[0].Children) != 0 {
+		t.Error("previously returned projection mutated by a later session write")
+	}
+	after := a.SessionTree()
+	if len(after) != 1 || len(after[0].Children) != 1 {
+		t.Errorf("fresh projection missing the new entry: %+v", after)
+	}
+}
+
+func TestDescribeEntryUnknown(t *testing.T) {
+	kind, role, text := describeEntry(struct{ Nonsense int }{})
+	if kind != EntryKindUnknown {
+		t.Errorf("Kind = %q, want %q", kind, EntryKindUnknown)
+	}
+	if role != "" || text != "" {
+		t.Errorf("role/text = %q/%q, want empty", role, text)
+	}
+}
+
+// --------------------------------------------------------------------------
+// SessionHistory
+// --------------------------------------------------------------------------
+
+func TestSessionHistoryNoSession(t *testing.T) {
+	a := New(Options{}, nil)
+	if h := a.SessionHistory(); h != nil {
+		t.Fatalf("expected nil history without a session, got %+v", h)
+	}
+}
+
+func TestSessionHistoryReturnsBranchMessages(t *testing.T) {
+	a, tm := newSessionApp(t)
+
+	appendUserMessage(t, tm, "hello")
+	appendMessage(t, tm, message.RoleAssistant, "hi there")
+	// Non-message entries on the branch must be skipped, not returned.
+	if _, err := tm.AppendModelChange("anthropic", "claude-sonnet-4-5"); err != nil {
+		t.Fatalf("AppendModelChange: %v", err)
+	}
+
+	history := a.SessionHistory()
+	if len(history) != 2 {
+		t.Fatalf("len(history) = %d, want 2: %+v", len(history), history)
+	}
+	if history[0].Role != message.RoleUser || history[0].Content() != "hello" {
+		t.Errorf("history[0] = %v/%q, want user/hello", history[0].Role, history[0].Content())
+	}
+	if history[1].Role != message.RoleAssistant || history[1].Content() != "hi there" {
+		t.Errorf("history[1] = %v/%q, want assistant/hi there", history[1].Role, history[1].Content())
+	}
+}
+
+// --------------------------------------------------------------------------
+// Mutations
+// --------------------------------------------------------------------------
+
+func TestSessionMutationsWithoutSession(t *testing.T) {
+	a := New(Options{}, nil)
+
+	if err := a.SetSessionName("nope"); !errors.Is(err, ErrNoSession) {
+		t.Errorf("SetSessionName err = %v, want ErrNoSession", err)
+	}
+	if err := a.NewSession(t.TempDir()); !errors.Is(err, ErrNoSession) {
+		t.Errorf("NewSession err = %v, want ErrNoSession", err)
+	}
+	if err := a.ForkSession(t.TempDir(), "abc"); !errors.Is(err, ErrNoSession) {
+		t.Errorf("ForkSession err = %v, want ErrNoSession", err)
+	}
+	if _, err := a.SessionSystemPromptEntry("prompt", "model"); !errors.Is(err, ErrNoSession) {
+		t.Errorf("SessionSystemPromptEntry err = %v, want ErrNoSession", err)
+	}
+}
+
+func TestSetSessionName(t *testing.T) {
+	a, _ := newSessionApp(t)
+
+	if err := a.SetSessionName("renamed"); err != nil {
+		t.Fatalf("SetSessionName: %v", err)
+	}
+	snap, ok := a.SessionSnapshot()
+	if !ok {
+		t.Fatal("expected an active session")
+	}
+	if snap.Name != "renamed" {
+		t.Errorf("Name = %q, want %q", snap.Name, "renamed")
+	}
+}
+
+func TestSessionSystemPromptEntryUsesSessionModel(t *testing.T) {
+	a, tm := newSessionApp(t)
+	appendUserMessage(t, tm, "hello")
+	if _, err := tm.AppendModelChange("anthropic", "claude-sonnet-4-5"); err != nil {
+		t.Fatalf("AppendModelChange: %v", err)
+	}
+
+	data, err := a.SessionSystemPromptEntry("be helpful", "fallback-model")
+	if err != nil {
+		t.Fatalf("SessionSystemPromptEntry: %v", err)
+	}
+
+	entry, err := session.UnmarshalEntry(data)
+	if err != nil {
+		t.Fatalf("UnmarshalEntry: %v", err)
+	}
+	sp, ok := entry.(*session.SystemPromptEntry)
+	if !ok {
+		t.Fatalf("got %T, want *session.SystemPromptEntry", entry)
+	}
+	if sp.Content != "be helpful" {
+		t.Errorf("Content = %q, want %q", sp.Content, "be helpful")
+	}
+	if sp.Model != "claude-sonnet-4-5" {
+		t.Errorf("Model = %q, want the session's model", sp.Model)
+	}
+	if sp.Provider != "anthropic" {
+		t.Errorf("Provider = %q, want %q", sp.Provider, "anthropic")
+	}
+}
+
+func TestSessionSystemPromptEntryFallsBackToGivenModel(t *testing.T) {
+	a, tm := newSessionApp(t)
+	appendUserMessage(t, tm, "hello")
+
+	data, err := a.SessionSystemPromptEntry("be helpful", "fallback-model")
+	if err != nil {
+		t.Fatalf("SessionSystemPromptEntry: %v", err)
+	}
+	entry, err := session.UnmarshalEntry(data)
+	if err != nil {
+		t.Fatalf("UnmarshalEntry: %v", err)
+	}
+	sp := entry.(*session.SystemPromptEntry)
+	if sp.Model != "fallback-model" {
+		t.Errorf("Model = %q, want the fallback when the session records none", sp.Model)
+	}
+}
diff --git a/internal/session/entry.go b/internal/session/entry.go
index a19b92a3..b540a0e5 100644
--- a/internal/session/entry.go
+++ b/internal/session/entry.go
@@ -5,6 +5,7 @@ import (
 	"encoding/hex"
 	"encoding/json"
 	"fmt"
+	"strings"
 	"time"
 
 	"github.com/mark3labs/kit/internal/message"
@@ -333,6 +334,32 @@ func UnmarshalEntry(data []byte) (any, error) {
 	}
 }
 
+// Text returns the plain-text content of the message, concatenating every
+// text part with newlines. Non-text parts (tool calls, tool results, images,
+// reasoning) are ignored, so a tool-only message yields an empty string.
+//
+// This is a cheap projection intended for previews and list rendering: it
+// scans the type-tagged parts directly instead of going through ToMessage,
+// and it never fails — malformed parts simply yield an empty string.
+func (e *MessageEntry) Text() string {
+	var parts []struct {
+		Type string `json:"type"`
+		Data struct {
+			Text string `json:"text"`
+		} `json:"data"`
+	}
+	if err := json.Unmarshal(e.Parts, &parts); err != nil {
+		return ""
+	}
+	var texts []string
+	for _, p := range parts {
+		if p.Type == "text" && p.Data.Text != "" {
+			texts = append(texts, p.Data.Text)
+		}
+	}
+	return strings.Join(texts, "\n")
+}
+
 // ToMessage converts a MessageEntry back to a message.Message by
 // unmarshaling the type-tagged parts.
 func (e *MessageEntry) ToMessage() (message.Message, error) {
diff --git a/internal/ui/core/events.go b/internal/ui/core/events.go
index 6c80d746..0c7808fc 100644
--- a/internal/ui/core/events.go
+++ b/internal/ui/core/events.go
@@ -36,8 +36,10 @@ type CtrlCResetMsg struct{}
 type TreeNodeSelectedMsg struct {
 	// ID is the entry ID of the selected node.
 	ID string
-	// Entry is the underlying entry object.
-	Entry any
+	// ParentID is the entry ID of the selected node's parent, empty for
+	// root entries. Forking a user message branches from the parent so the
+	// message itself can be edited and resubmitted.
+	ParentID string
 	// IsUser is true if the selected entry is a user message.
 	IsUser bool
 	// UserText is the user message text (only set when IsUser is true).
diff --git a/internal/ui/model.go b/internal/ui/model.go
index 35c06145..139eb60a 100644
--- a/internal/ui/model.go
+++ b/internal/ui/model.go
@@ -22,7 +22,6 @@ import (
 	"github.com/mark3labs/kit/internal/message"
 	"github.com/mark3labs/kit/internal/models"
 	"github.com/mark3labs/kit/internal/prompts"
-	"github.com/mark3labs/kit/internal/session"
 	"github.com/mark3labs/kit/internal/ui/clipboard"
 	"github.com/mark3labs/kit/internal/ui/commands"
 	uicore "github.com/mark3labs/kit/internal/ui/core"
@@ -102,12 +101,31 @@ type AppController interface {
 	// error synchronously if compaction cannot be started (e.g. agent is busy).
 	// customInstructions is optional text appended to the summary prompt.
 	CompactConversation(customInstructions string) error
-	// GetTreeSession returns the tree session manager, or nil if tree sessions
-	// are not enabled. Used by slash commands like /tree, /fork, /session.
-	GetTreeSession() *session.TreeManager
-	// SwitchTreeSession replaces the active tree session with a new one,
-	// closing the old session. Used by /new to create a completely fresh session.
-	SwitchTreeSession(ts *session.TreeManager)
+	// SessionSnapshot returns a value snapshot of the active tree session's
+	// metadata (id, name, file path, counts, current leaf). ok is false when
+	// no tree session is active. Callers re-read it whenever they need
+	// current values rather than caching it.
+	SessionSnapshot() (app.SessionSnapshot, bool)
+	// SessionTree returns the session's entry tree projected into value
+	// nodes, or nil when no tree session is active. Used by /tree and /fork.
+	SessionTree() []app.TreeNodeView
+	// SessionHistory returns the conversation messages on the session's
+	// current branch, oldest first. Used to replay the transcript after
+	// resuming or forking a session.
+	SessionHistory() []message.Message
+	// SetSessionName sets the session's display name. Used by /name.
+	SetSessionName(name string) error
+	// NewSession replaces the active session with a brand new one created in
+	// cwd, closing the old one. Used by /new.
+	NewSession(cwd string) error
+	// ForkSession creates a new session in cwd holding the history up to
+	// targetID and switches to it. Used by /fork and tree-node selection.
+	ForkSession(cwd, targetID string) error
+	// SessionSystemPromptEntry marshals a system-prompt entry describing the
+	// given prompt plus the session's current model/provider, for embedding
+	// in exported and shared session files. fallbackModelID is used when the
+	// session records no model of its own.
+	SessionSystemPromptEntry(systemPrompt, fallbackModelID string) ([]byte, error)
 	// SendUIMessage re-injects a UI-internal message into the program's Update
 	// loop asynchronously. Safe to call from any goroutine. Used by extension
 	// command goroutines (and other async UI work) to deliver results back to
@@ -1520,17 +1538,13 @@ func (m *AppModel) update(msg tea.Msg) (tea.Model, tea.Cmd) {
 	// ── Tree selector events ─────────────────────────────────────────────────
 	case uicore.TreeNodeSelectedMsg:
 		// User selected a node in the tree. Branch to it and return to input.
-		if ts := m.appCtrl.GetTreeSession(); ts != nil {
+		if _, ok := m.appCtrl.SessionSnapshot(); ok {
 			// For user messages: branch to parent (so user can resubmit).
 			// For other entries: branch directly to the selected entry.
 			targetID := msg.ID
 			if msg.IsUser {
 				// Branch to parent of user message, place text in editor.
-				if node := ts.GetEntry(msg.ID); node != nil {
-					if me, ok := node.(*session.MessageEntry); ok {
-						targetID = me.ParentID
-					}
-				}
+				targetID = msg.ParentID
 			}
 
 			// Emit before-fork event in a goroutine so that extension handlers
@@ -5115,17 +5129,17 @@ func (m *AppModel) handleThinkingCommand(args string) tea.Cmd {
 
 // handleTreeCommand opens the tree selector overlay.
 func (m *AppModel) handleTreeCommand() tea.Cmd {
-	ts := m.appCtrl.GetTreeSession()
-	if ts == nil {
+	snap, ok := m.appCtrl.SessionSnapshot()
+	if !ok {
 		m.printSystemMessage("No tree session active. Start with `--continue` or `--resume` to enable tree sessions.")
 		return nil
 	}
-	if ts.EntryCount() == 0 {
+	if snap.EntryCount == 0 {
 		m.printSystemMessage("No entries in session yet.")
 		return nil
 	}
 
-	m.treeSelector = NewTreeSelector(ts, m.width, m.height)
+	m.treeSelector = NewTreeSelector(m.appCtrl.SessionTree(), snap.LeafID, m.width, m.height)
 	m.state = stateTreeSelector
 	return nil
 }
@@ -5135,18 +5149,18 @@ func (m *AppModel) handleTreeCommand() tea.Cmd {
 // Unlike /tree which shows the full tree, /fork shows only user messages
 // (matching Pi's behavior) and creates a new session file when a message is selected.
 func (m *AppModel) handleForkCommand() tea.Cmd {
-	ts := m.appCtrl.GetTreeSession()
-	if ts == nil {
+	snap, ok := m.appCtrl.SessionSnapshot()
+	if !ok {
 		m.printSystemMessage("No tree session active. Start with `--continue` or `--resume` to enable tree sessions.")
 		return nil
 	}
-	if ts.EntryCount() == 0 {
+	if snap.EntryCount == 0 {
 		m.printSystemMessage("No entries to fork from.")
 		return nil
 	}
 
 	// Use the fork-specific selector that shows only user messages.
-	m.treeSelector = NewTreeSelectorForFork(ts, m.width, m.height)
+	m.treeSelector = NewTreeSelectorForFork(m.appCtrl.SessionTree(), snap.LeafID, m.width, m.height)
 	m.state = stateTreeSelector
 	return nil
 }
@@ -5185,8 +5199,7 @@ func (m *AppModel) handleNewCommand(initialPrompt string) tea.Cmd {
 // context from the previous conversation. If initialPrompt is non-empty it
 // is submitted as the first user turn (with @file expansion).
 func (m *AppModel) performNewSession(initialPrompt string) tea.Cmd {
-	ts := m.appCtrl.GetTreeSession()
-	if ts == nil {
+	if _, ok := m.appCtrl.SessionSnapshot(); !ok {
 		// No tree session — just clear messages.
 		if m.appCtrl != nil {
 			m.appCtrl.ClearMessages()
@@ -5203,16 +5216,14 @@ func (m *AppModel) performNewSession(initialPrompt string) tea.Cmd {
 		return cmd
 	}
 
-	// Create a brand new session file (Pi-style /new behavior)
-	newTs, err := session.CreateTreeSession(m.cwd)
-	if err != nil {
+	// Create a brand new session file (Pi-style /new behavior) and switch to
+	// it, closing the old one.
+	if err := m.appCtrl.NewSession(m.cwd); err != nil {
 		m.printSystemMessage(fmt.Sprintf("Failed to create new session: %v", err))
 		m.signalNewSessionResult(fmt.Errorf("create new session: %w", err))
 		return nil
 	}
 
-	// Switch to the new session, closing the old one
-	m.appCtrl.SwitchTreeSession(newTs)
 	// Reset usage statistics for the new session
 	if m.usageTracker != nil {
 		m.usageTracker.Reset()
@@ -5290,23 +5301,19 @@ func (m *AppModel) submitInitialPrompt(prompt string) tea.Cmd {
 // Called either directly (when no before-hook exists) or after the async
 // before-fork hook completes.
 func (m *AppModel) performFork(targetID string, isUser bool, userText string) tea.Cmd {
-	ts := m.appCtrl.GetTreeSession()
-	if ts == nil {
+	if _, ok := m.appCtrl.SessionSnapshot(); !ok {
 		m.printSystemMessage("No tree session active.")
 		return nil
 	}
 
-	// Create a new session by forking from the target entry.
-	// This creates a new session file with the history up to the target point.
-	newTs, err := ts.ForkToNewSession(m.cwd, targetID)
-	if err != nil {
+	// Create a new session by forking from the target entry. This writes a
+	// new session file with the history up to the target point and switches
+	// to it.
+	if err := m.appCtrl.ForkSession(m.cwd, targetID); err != nil {
 		m.printSystemMessage(fmt.Sprintf("Failed to fork session: %v", err))
 		return nil
 	}
 
-	// Switch to the new forked session.
-	m.appCtrl.SwitchTreeSession(newTs)
-
 	// Reset usage statistics for the new session.
 	if m.usageTracker != nil {
 		m.usageTracker.Reset()
@@ -5333,17 +5340,16 @@ func (m *AppModel) performFork(targetID string, isUser bool, userText string) te
 //
 //	/name             — shows the current name.
 func (m *AppModel) handleNameCommand(args string) tea.Cmd {
-	ts := m.appCtrl.GetTreeSession()
-	if ts == nil {
+	snap, ok := m.appCtrl.SessionSnapshot()
+	if !ok {
 		m.printSystemMessage("No tree session active.")
 		return nil
 	}
 
 	if args == "" {
 		// No argument — show current name.
-		currentName := ts.GetSessionName()
-		if currentName != "" {
-			m.printSystemMessage(fmt.Sprintf("Session name: %q\nTo rename: `/name <new name>`", currentName))
+		if snap.Name != "" {
+			m.printSystemMessage(fmt.Sprintf("Session name: %q\nTo rename: `/name <new name>`", snap.Name))
 		} else {
 			m.printSystemMessage("Session has no name. Set one with: `/name <new name>`")
 		}
@@ -5351,7 +5357,7 @@ func (m *AppModel) handleNameCommand(args string) tea.Cmd {
 	}
 
 	// Set the session name.
-	if _, err := ts.AppendSessionInfo(args); err != nil {
+	if err := m.appCtrl.SetSessionName(args); err != nil {
 		m.printSystemMessage(fmt.Sprintf("Failed to set session name: %v", err))
 		return nil
 	}
@@ -5584,13 +5590,13 @@ func (m *AppModel) handleEditCommand(args string) tea.Cmd {
 //
 //	/export path.jsonl — copies to the specified path.
 func (m *AppModel) handleExportCommand(args string) tea.Cmd {
-	ts := m.appCtrl.GetTreeSession()
-	if ts == nil {
+	snap, ok := m.appCtrl.SessionSnapshot()
+	if !ok {
 		m.printSystemMessage("No tree session active.")
 		return nil
 	}
 
-	srcPath := ts.GetFilePath()
+	srcPath := snap.FilePath
 	if srcPath == "" {
 		m.printSystemMessage("Session is in-memory (not persisted). Nothing to export.")
 		return nil
@@ -5600,17 +5606,12 @@ func (m *AppModel) handleExportCommand(args string) tea.Cmd {
 	dstPath := args
 	if dstPath == "" {
 		// Generate a name based on session name or ID.
-		name := ts.GetSessionName()
+		name := snap.Name
 		if name == "" {
-			name = ts.GetSessionID()[:12]
+			name = shortSessionID(snap.ID)
 		}
 		// Sanitize for filename.
-		name = strings.Map(func(r rune) rune {
-			if r == '/' || r == '\\' || r == ':' || r == ' ' {
-				return '_'
-			}
-			return r
-		}, name)
+		name = sanitizeFileName(name)
 		dstPath = fmt.Sprintf("session_%s.jsonl", name)
 	}
 
@@ -5634,13 +5635,13 @@ func (m *AppModel) handleExportCommand(args string) tea.Cmd {
 // a shareable viewer URL. Requires the GitHub CLI (gh) to be installed and
 // authenticated.
 func (m *AppModel) handleShareCommand() tea.Cmd {
-	ts := m.appCtrl.GetTreeSession()
-	if ts == nil {
+	snap, ok := m.appCtrl.SessionSnapshot()
+	if !ok {
 		m.printSystemMessage("No tree session active.")
 		return nil
 	}
 
-	srcPath := ts.GetFilePath()
+	srcPath := snap.FilePath
 	if srcPath == "" {
 		m.printSystemMessage("Session is in-memory (not persisted). Nothing to share.")
 		return nil
@@ -5666,33 +5667,23 @@ func (m *AppModel) handleShareCommand() tea.Cmd {
 		return nil
 	}
 
-	// Capture the current system prompt and model info.
-	systemPrompt := viper.GetString("system-prompt")
-	_, provider, modelID := ts.BuildContext()
-	if modelID == "" {
-		// Fallback to viper if no model change recorded in session
-		modelID = viper.GetString("model")
-	}
-
-	// Create a SystemPromptEntry with both prompt and model info.
-	sysPromptEntry := session.NewSystemPromptEntry(systemPrompt, modelID, provider)
-	sysPromptJSON, err := session.MarshalEntry(sysPromptEntry)
+	// Capture the current system prompt and model info as a system-prompt
+	// entry so the shared file records the context the conversation ran under.
+	sysPromptJSON, err := m.appCtrl.SessionSystemPromptEntry(
+		viper.GetString("system-prompt"),
+		viper.GetString("model"),
+	)
 	if err != nil {
 		m.printSystemMessage(fmt.Sprintf("Failed to marshal system prompt: %v", err))
 		return nil
 	}
 
-	name := ts.GetSessionName()
+	name := snap.Name
 	if name == "" {
 		name = "session"
 	}
 	// Sanitize for filename.
-	name = strings.Map(func(r rune) rune {
-		if r == '/' || r == '\\' || r == ':' || r == ' ' {
-			return '_'
-		}
-		return r
-	}, name)
+	name = sanitizeFileName(name)
 
 	tmpPath, err := buildShareFile(name, data, sysPromptJSON)
 	if err != nil {
@@ -5820,13 +5811,8 @@ func (m *AppModel) handleResumeCommand() tea.Cmd {
 // This gives the user visual context of the conversation when resuming or
 // importing a session. Call this after switchSession succeeds.
 func (m *AppModel) renderSessionHistory() {
-	ts := m.appCtrl.GetTreeSession()
-	if ts == nil {
-		return
-	}
-
-	branch := ts.GetBranch("")
-	if len(branch) == 0 {
+	history := m.appCtrl.SessionHistory()
+	if len(history) == 0 {
 		return
 	}
 
@@ -5840,16 +5826,8 @@ func (m *AppModel) renderSessionHistory() {
 		Args string
 	}
 	toolCallMap := make(map[string]toolCallInfo)
-	for _, entry := range branch {
-		me, ok := entry.(*session.MessageEntry)
-		if !ok {
-			continue
-		}
-		if me.Role != "assistant" {
-			continue
-		}
-		msg, err := me.ToMessage()
-		if err != nil {
+	for _, msg := range history {
+		if msg.Role != message.RoleAssistant {
 			continue
 		}
 		for _, tc := range msg.ToolCalls() {
@@ -5858,16 +5836,7 @@ func (m *AppModel) renderSessionHistory() {
 	}
 
 	// Second pass: create MessageItems for each message in order.
-	for _, entry := range branch {
-		me, ok := entry.(*session.MessageEntry)
-		if !ok {
-			continue
-		}
-		msg, err := me.ToMessage()
-		if err != nil {
-			continue
-		}
-
+	for _, msg := range history {
 		switch msg.Role {
 		case message.RoleUser:
 			text := strings.TrimSpace(msg.Content())
@@ -5940,15 +5909,35 @@ func (m *AppModel) renderSessionHistory() {
 	m.pendingGotoBottom = true
 }
 
+// sanitizeFileName replaces path separators and other characters that are
+// awkward in file names with underscores, so a user-chosen session name can
+// be embedded in an export/share filename safely.
+func sanitizeFileName(name string) string {
+	return strings.Map(func(r rune) rune {
+		if r == '/' || r == '\\' || r == ':' || r == ' ' {
+			return '_'
+		}
+		return r
+	}, name)
+}
+
+// shortSessionID returns a filename-friendly prefix of a session ID, used as
+// a fallback export name for unnamed sessions.
+func shortSessionID(id string) string {
+	if len(id) > 12 {
+		return id[:12]
+	}
+	return id
+}
+
 // handleSessionInfoCommand shows session statistics.
 func (m *AppModel) handleSessionInfoCommand() tea.Cmd {
-	ts := m.appCtrl.GetTreeSession()
-	if ts == nil {
+	snap, ok := m.appCtrl.SessionSnapshot()
+	if !ok {
 		m.printSystemMessage("No tree session active.")
 		return nil
 	}
 
-	header := ts.GetHeader()
 	info := fmt.Sprintf("## Session Info\n\n"+
 		"- **ID:** `%s`\n"+
 		"- **File:** `%s`\n"+
@@ -5957,17 +5946,17 @@ func (m *AppModel) handleSessionInfoCommand() tea.Cmd {
 		"- **Entries:** %d\n"+
 		"- **Messages:** %d\n"+
 		"- **Current Leaf:** `%s`\n",
-		header.ID,
-		ts.GetFilePath(),
-		header.Cwd,
-		header.Timestamp.Format(time.RFC3339),
-		ts.EntryCount(),
-		ts.MessageCount(),
-		ts.GetLeafID(),
+		snap.ID,
+		snap.FilePath,
+		snap.Cwd,
+		snap.Created.Format(time.RFC3339),
+		snap.EntryCount,
+		snap.MessageCount,
+		snap.LeafID,
 	)
 
-	if name := ts.GetSessionName(); name != "" {
-		info += fmt.Sprintf("- **Name:** %s\n", name)
+	if snap.Name != "" {
+		info += fmt.Sprintf("- **Name:** %s\n", snap.Name)
 	}
 
 	m.printSystemMessage(info)
diff --git a/internal/ui/model_test.go b/internal/ui/model_test.go
index a96349c2..315d09b0 100644
--- a/internal/ui/model_test.go
+++ b/internal/ui/model_test.go
@@ -8,7 +8,7 @@ import (
 
 	tea "charm.land/bubbletea/v2"
 	"github.com/mark3labs/kit/internal/app"
-	"github.com/mark3labs/kit/internal/session"
+	"github.com/mark3labs/kit/internal/message"
 	"github.com/mark3labs/kit/internal/ui/core"
 	kit "github.com/mark3labs/kit/pkg/kit"
 )
@@ -57,12 +57,32 @@ func (s *stubAppController) CompactConversation(_ string) error {
 	return nil
 }
 
-func (s *stubAppController) GetTreeSession() *session.TreeManager {
+func (s *stubAppController) SessionSnapshot() (app.SessionSnapshot, bool) {
+	return app.SessionSnapshot{}, false
+}
+
+func (s *stubAppController) SessionTree() []app.TreeNodeView {
 	return nil
 }
 
-func (s *stubAppController) SwitchTreeSession(_ *session.TreeManager) {
-	// no-op in tests
+func (s *stubAppController) SessionHistory() []message.Message {
+	return nil
+}
+
+func (s *stubAppController) SetSessionName(_ string) error {
+	return app.ErrNoSession
+}
+
+func (s *stubAppController) NewSession(_ string) error {
+	return app.ErrNoSession
+}
+
+func (s *stubAppController) ForkSession(_, _ string) error {
+	return app.ErrNoSession
+}
+
+func (s *stubAppController) SessionSystemPromptEntry(_, _ string) ([]byte, error) {
+	return nil, app.ErrNoSession
 }
 
 func (s *stubAppController) SendUIMessage(_ tea.Msg) {
diff --git a/internal/ui/tree_selector.go b/internal/ui/tree_selector.go
index f5276876..7fb79f66 100644
--- a/internal/ui/tree_selector.go
+++ b/internal/ui/tree_selector.go
@@ -1,7 +1,6 @@
 package ui
 
 import (
-	"encoding/json"
 	"fmt"
 	"strings"
 
@@ -9,7 +8,7 @@ import (
 	tea "charm.land/bubbletea/v2"
 	"charm.land/lipgloss/v2"
 
-	"github.com/mark3labs/kit/internal/session"
+	"github.com/mark3labs/kit/internal/app"
 	"github.com/mark3labs/kit/internal/ui/core"
 )
 
@@ -43,8 +42,8 @@ func (m TreeFilterMode) String() string {
 
 // FlatNode is a tree entry flattened for list rendering with indentation info.
 type FlatNode struct {
-	Entry    any    // the underlying entry
-	ID       string // entry ID
+	Node     app.TreeNodeView // value projection of the entry
+	ID       string           // entry ID
 	ParentID string
 	Depth    int    // indentation level
 	IsLast   bool   // last child at this depth
@@ -59,8 +58,8 @@ type FlatNode struct {
 // filtered node list and a custom RenderItem that draws each tree node with
 // its indentation prefix and role colors.
 type TreeSelectorComponent struct {
-	tm         *session.TreeManager
-	flatNodes  []FlatNode // visible nodes (matches popup.Items() 1:1)
+	roots      []app.TreeNodeView // value projection of the session tree
+	flatNodes  []FlatNode         // visible nodes (matches popup.Items() 1:1)
 	filter     TreeFilterMode
 	leafID     string // real leaf for "active" marker
 	popup      *PopupList
@@ -71,12 +70,13 @@ type TreeSelectorComponent struct {
 	cancelled  bool
 }
 
-// NewTreeSelector creates a tree selector from a TreeManager.
-func NewTreeSelector(tm *session.TreeManager, width, height int) *TreeSelectorComponent {
+// NewTreeSelector creates a tree selector over a session tree snapshot.
+// roots is the projected entry tree and leafID marks the active branch tip.
+func NewTreeSelector(roots []app.TreeNodeView, leafID string, width, height int) *TreeSelectorComponent {
 	ts := &TreeSelectorComponent{
-		tm:     tm,
+		roots:  roots,
 		filter: TreeFilterDefault,
-		leafID: tm.GetLeafID(),
+		leafID: leafID,
 		width:  width,
 		height: height,
 		active: true,
@@ -95,11 +95,11 @@ func NewTreeSelector(tm *session.TreeManager, width, height int) *TreeSelectorCo
 
 // NewTreeSelectorForFork creates a tree selector for the /fork command.
 // It shows only user messages (flat list) matching Pi's fork behavior.
-func NewTreeSelectorForFork(tm *session.TreeManager, width, height int) *TreeSelectorComponent {
+func NewTreeSelectorForFork(roots []app.TreeNodeView, leafID string, width, height int) *TreeSelectorComponent {
 	ts := &TreeSelectorComponent{
-		tm:     tm,
+		roots:  roots,
 		filter: TreeFilterUserOnly,
-		leafID: tm.GetLeafID(),
+		leafID: leafID,
 		width:  width,
 		height: height,
 		active: true,
@@ -108,7 +108,7 @@ func NewTreeSelectorForFork(tm *session.TreeManager, width, height int) *TreeSel
 	ts.rebuild()
 	// Position cursor at the last user message before the leaf.
 	for i := len(ts.flatNodes) - 1; i >= 0; i-- {
-		if ts.isUserMessage(ts.flatNodes[i].Entry) {
+		if ts.flatNodes[i].Node.IsUserMessage() {
 			ts.popup.SetCursor(i)
 			break
 		}
@@ -191,9 +191,9 @@ func (ts *TreeSelectorComponent) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 				return ts, func() tea.Msg {
 					return core.TreeNodeSelectedMsg{
 						ID:       node.ID,
-						Entry:    node.Entry,
-						IsUser:   ts.isUserMessage(node.Entry),
-						UserText: ts.extractUserText(node.Entry),
+						ParentID: node.ParentID,
+						IsUser:   node.Node.IsUserMessage(),
+						UserText: userText(node.Node),
 					}
 				}
 			}
@@ -235,9 +235,8 @@ func (ts *TreeSelectorComponent) IsActive() bool {
 // with PopupItems. Called on initial load and whenever the filter changes.
 func (ts *TreeSelectorComponent) rebuild() {
 	ts.flatNodes = ts.flatNodes[:0]
-	tree := ts.tm.GetTree()
-	for i, root := range tree {
-		isLast := i == len(tree)-1
+	for i, root := range ts.roots {
+		isLast := i == len(ts.roots)-1
 		ts.flattenNode(root, 0, isLast, "")
 	}
 	ts.publishItems()
@@ -263,7 +262,7 @@ func (ts *TreeSelectorComponent) publishItems() {
 	items := make([]PopupItem, len(ts.flatNodes))
 	for i, n := range ts.flatNodes {
 		items[i] = PopupItem{
-			Label:  ts.entryDisplayText(n.Entry),
+			Label:  entryDisplayText(n.Node),
 			Active: n.ID == ts.leafID,
 			Meta:   n,
 		}
@@ -273,7 +272,7 @@ func (ts *TreeSelectorComponent) publishItems() {
 	ts.syncFlatNodes()
 }
 
-func (ts *TreeSelectorComponent) flattenNode(node *session.TreeNode, depth int, isLast bool, gutterPrefix string) {
+func (ts *TreeSelectorComponent) flattenNode(node app.TreeNodeView, depth int, isLast bool, gutterPrefix string) {
 	if !ts.passesFilter(node) {
 		// Still recurse into children in case they pass.
 		for i, child := range node.Children {
@@ -292,16 +291,14 @@ func (ts *TreeSelectorComponent) flattenNode(node *session.TreeNode, depth int,
 		prefix = gutterPrefix + "├─ "
 	}
 
-	label := ts.tm.GetLabel(node.ID)
-
 	ts.flatNodes = append(ts.flatNodes, FlatNode{
-		Entry:    node.Entry,
+		Node:     node,
 		ID:       node.ID,
 		ParentID: node.ParentID,
 		Depth:    depth,
 		IsLast:   isLast,
 		Prefix:   prefix,
-		Label:    label,
+		Label:    node.Label,
 	})
 
 	// Build gutter prefix for children.
@@ -320,39 +317,34 @@ func (ts *TreeSelectorComponent) flattenNode(node *session.TreeNode, depth int,
 	}
 }
 
-func (ts *TreeSelectorComponent) passesFilter(node *session.TreeNode) bool {
+func (ts *TreeSelectorComponent) passesFilter(node app.TreeNodeView) bool {
 	switch ts.filter {
 	case TreeFilterAll:
 		return true
 
 	case TreeFilterDefault:
 		// Hide settings entries.
-		switch node.Entry.(type) {
-		case *session.ModelChangeEntry, *session.LabelEntry, *session.SessionInfoEntry:
+		switch node.Kind {
+		case app.EntryKindModelChange, app.EntryKindLabel, app.EntryKindSessionInfo:
 			return false
 		}
 		// Hide tool messages unless they're the leaf.
-		if me, ok := node.Entry.(*session.MessageEntry); ok {
-			if me.Role == "tool" && node.ID != ts.leafID {
-				return false
-			}
+		if node.Kind == app.EntryKindMessage && node.Role == "tool" && node.ID != ts.leafID {
+			return false
 		}
 		return true
 
 	case TreeFilterNoTools:
-		if me, ok := node.Entry.(*session.MessageEntry); ok {
-			return me.Role != "tool"
+		if node.Kind == app.EntryKindMessage {
+			return node.Role != "tool"
 		}
 		return true
 
 	case TreeFilterUserOnly:
-		if me, ok := node.Entry.(*session.MessageEntry); ok {
-			return me.Role == "user"
-		}
-		return false
+		return node.IsUserMessage()
 
 	case TreeFilterLabelOnly:
-		return ts.tm.GetLabel(node.ID) != ""
+		return node.Label != ""
 
 	default:
 		return true
@@ -415,7 +407,7 @@ func (ts *TreeSelectorComponent) renderNode(item PopupItem, innerWidth int, isCu
 	// Reserve space for indicator(2) + prefix + right parts.
 	available := max(innerWidth-2-prefixW-rightW, 4)
 
-	text := ts.entryDisplayText(node.Entry)
+	text := entryDisplayText(node.Node)
 	text = truncateRunes(text, available)
 
 	// Selected row: emit raw text. The outer row style applies fg+bg in one
@@ -426,9 +418,9 @@ func (ts *TreeSelectorComponent) renderNode(item PopupItem, innerWidth int, isCu
 
 	// Role-based text color.
 	var textStyle lipgloss.Style
-	switch e := node.Entry.(type) {
-	case *session.MessageEntry:
-		switch e.Role {
+	switch node.Node.Kind {
+	case app.EntryKindMessage:
+		switch node.Node.Role {
 		case "user":
 			textStyle = lipgloss.NewStyle().Foreground(theme.Accent)
 		case "assistant":
@@ -436,9 +428,9 @@ func (ts *TreeSelectorComponent) renderNode(item PopupItem, innerWidth int, isCu
 		default:
 			textStyle = lipgloss.NewStyle().Foreground(theme.Muted)
 		}
-	case *session.BranchSummaryEntry:
+	case app.EntryKindBranchSummary:
 		textStyle = lipgloss.NewStyle().Foreground(theme.Warning).Italic(true)
-	case *session.CompactionEntry:
+	case app.EntryKindCompaction:
 		textStyle = lipgloss.NewStyle().Foreground(theme.Info).Italic(true)
 	default:
 		textStyle = lipgloss.NewStyle().Foreground(theme.Muted)
@@ -458,31 +450,33 @@ func (ts *TreeSelectorComponent) renderNode(item PopupItem, innerWidth int, isCu
 	return parts
 }
 
-func (ts *TreeSelectorComponent) entryDisplayText(entry any) string {
-	switch e := entry.(type) {
-	case *session.MessageEntry:
-		role := e.Role
-		text := collapseToLine(extractTextFromParts(e.Parts))
-		text = truncateRunes(text, 200)
+// entryDisplayText renders a one-line summary of a tree node. It switches on
+// the app-layer EntryKind rather than on session storage types, so new entry
+// kinds degrade to the "(unknown entry)" fallback instead of breaking the
+// build.
+func entryDisplayText(node app.TreeNodeView) string {
+	switch node.Kind {
+	case app.EntryKindMessage:
+		text := truncateRunes(collapseToLine(node.Text), 200)
 		if text == "" {
 			text = "(tool interaction)"
 		}
-		return fmt.Sprintf("%s: %s", role, text)
+		return fmt.Sprintf("%s: %s", node.Role, text)
 
-	case *session.ModelChangeEntry:
-		return fmt.Sprintf("model: %s/%s", e.Provider, e.ModelID)
+	case app.EntryKindModelChange:
+		return fmt.Sprintf("model: %s", node.Text)
 
-	case *session.BranchSummaryEntry:
-		return fmt.Sprintf("branch summary: %s", truncateRunes(collapseToLine(e.Summary), 200))
+	case app.EntryKindBranchSummary:
+		return fmt.Sprintf("branch summary: %s", truncateRunes(collapseToLine(node.Text), 200))
 
-	case *session.CompactionEntry:
-		return fmt.Sprintf("compaction: %s", truncateRunes(collapseToLine(e.Summary), 200))
+	case app.EntryKindCompaction:
+		return fmt.Sprintf("compaction: %s", truncateRunes(collapseToLine(node.Text), 200))
 
-	case *session.LabelEntry:
-		return fmt.Sprintf("label: %s", e.Label)
+	case app.EntryKindLabel:
+		return fmt.Sprintf("label: %s", node.Text)
 
-	case *session.SessionInfoEntry:
-		return fmt.Sprintf("name: %s", e.Name)
+	case app.EntryKindSessionInfo:
+		return fmt.Sprintf("name: %s", node.Text)
 
 	default:
 		return "(unknown entry)"
@@ -496,37 +490,11 @@ func collapseToLine(s string) string {
 	return strings.Join(strings.Fields(s), " ")
 }
 
-func (ts *TreeSelectorComponent) isUserMessage(entry any) bool {
-	if me, ok := entry.(*session.MessageEntry); ok {
-		return me.Role == "user"
-	}
-	return false
-}
-
-func (ts *TreeSelectorComponent) extractUserText(entry any) string {
-	if me, ok := entry.(*session.MessageEntry); ok && me.Role == "user" {
-		return extractTextFromParts(me.Parts)
+// userText returns the full, untruncated text of a user message so it can be
+// placed back into the editor on fork. Empty for every other node.
+func userText(node app.TreeNodeView) string {
+	if node.IsUserMessage() {
+		return node.Text
 	}
 	return ""
 }
-
-// extractTextFromParts extracts text content from type-tagged parts JSON.
-func extractTextFromParts(partsJSON []byte) string {
-	// Quick extraction without full unmarshal.
-	var parts []struct {
-		Type string `json:"type"`
-		Data struct {
-			Text string `json:"text"`
-		} `json:"data"`
-	}
-	if err := json.Unmarshal(partsJSON, &parts); err != nil {
-		return ""
-	}
-	var texts []string
-	for _, p := range parts {
-		if p.Type == "text" && p.Data.Text != "" {
-			texts = append(texts, p.Data.Text)
-		}
-	}
-	return strings.Join(texts, "\n")
-}
diff --git a/internal/ui/tree_selector_test.go b/internal/ui/tree_selector_test.go
new file mode 100644
index 00000000..12132320
--- /dev/null
+++ b/internal/ui/tree_selector_test.go
@@ -0,0 +1,231 @@
+package ui
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/mark3labs/kit/internal/app"
+)
+
+// msgNode builds a message node for filter/render tests.
+func msgNode(id, role, text string) app.TreeNodeView {
+	return app.TreeNodeView{ID: id, Kind: app.EntryKindMessage, Role: role, Text: text}
+}
+
+func TestEntryDisplayText(t *testing.T) {
+	tests := []struct {
+		name string
+		node app.TreeNodeView
+		want string
+	}{
+		{
+			name: "user message",
+			node: msgNode("1", "user", "hello there"),
+			want: "user: hello there",
+		},
+		{
+			name: "multi-line message collapses to one line",
+			node: msgNode("1", "assistant", "line one\n\nline\ttwo"),
+			want: "assistant: line one line two",
+		},
+		{
+			name: "textless message is labelled as a tool interaction",
+			node: msgNode("1", "assistant", ""),
+			want: "assistant: (tool interaction)",
+		},
+		{
+			name: "model change",
+			node: app.TreeNodeView{Kind: app.EntryKindModelChange, Text: "anthropic/claude-sonnet-4-5"},
+			want: "model: anthropic/claude-sonnet-4-5",
+		},
+		{
+			name: "branch summary",
+			node: app.TreeNodeView{Kind: app.EntryKindBranchSummary, Text: "explored the parser"},
+			want: "branch summary: explored the parser",
+		},
+		{
+			name: "compaction",
+			node: app.TreeNodeView{Kind: app.EntryKindCompaction, Text: "earlier work"},
+			want: "compaction: earlier work",
+		},
+		{
+			name: "label",
+			node: app.TreeNodeView{Kind: app.EntryKindLabel, Text: "checkpoint"},
+			want: "label: checkpoint",
+		},
+		{
+			name: "session info",
+			node: app.TreeNodeView{Kind: app.EntryKindSessionInfo, Text: "my session"},
+			want: "name: my session",
+		},
+		{
+			name: "unknown kinds fall back rather than panic",
+			node: app.TreeNodeView{Kind: app.EntryKindUnknown},
+			want: "(unknown entry)",
+		},
+		{
+			name: "extension data has no preview",
+			node: app.TreeNodeView{Kind: app.EntryKindExtensionData},
+			want: "(unknown entry)",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := entryDisplayText(tt.node); got != tt.want {
+				t.Errorf("entryDisplayText() = %q, want %q", got, tt.want)
+			}
+		})
+	}
+}
+
+func TestEntryDisplayTextTruncatesLongMessages(t *testing.T) {
+	long := strings.Repeat("a", 500)
+	got := entryDisplayText(msgNode("1", "user", long))
+
+	// "user: " + 200 runes (the last of which is the ellipsis).
+	if want := len([]rune("user: ")) + 200; len([]rune(got)) != want {
+		t.Errorf("len = %d runes, want %d", len([]rune(got)), want)
+	}
+	if !strings.HasSuffix(got, "…") {
+		t.Errorf("expected a truncation ellipsis, got %q", got)
+	}
+}
+
+func TestUserText(t *testing.T) {
+	// The editor is repopulated from this, so it must not be truncated.
+	long := strings.Repeat("b", 500)
+	if got := userText(msgNode("1", "user", long)); got != long {
+		t.Errorf("user text was altered: len %d, want %d", len(got), len(long))
+	}
+	if got := userText(msgNode("1", "assistant", "nope")); got != "" {
+		t.Errorf("userText(assistant) = %q, want empty", got)
+	}
+	if got := userText(app.TreeNodeView{Kind: app.EntryKindCompaction, Text: "nope"}); got != "" {
+		t.Errorf("userText(compaction) = %q, want empty", got)
+	}
+}
+
+func TestPassesFilter(t *testing.T) {
+	const leafID = "leaf"
+
+	nodes := map[string]app.TreeNodeView{
+		"user":        msgNode("user", "user", "question"),
+		"assistant":   msgNode("assistant", "assistant", "answer"),
+		"tool":        msgNode("tool", "tool", ""),
+		"leafTool":    msgNode(leafID, "tool", ""),
+		"modelChange": {ID: "modelChange", Kind: app.EntryKindModelChange},
+		"label":       {ID: "label", Kind: app.EntryKindLabel},
+		"sessionInfo": {ID: "sessionInfo", Kind: app.EntryKindSessionInfo},
+		"compaction":  {ID: "compaction", Kind: app.EntryKindCompaction},
+		"labelled":    {ID: "labelled", Kind: app.EntryKindCompaction, Label: "pinned"},
+	}
+
+	// want[filter][nodeKey] = should the node be visible
+	want := map[TreeFilterMode]map[string]bool{
+		TreeFilterAll: {
+			"user": true, "assistant": true, "tool": true, "leafTool": true,
+			"modelChange": true, "label": true, "sessionInfo": true,
+			"compaction": true, "labelled": true,
+		},
+		TreeFilterDefault: {
+			// Settings entries are noise; tool messages are hidden unless
+			// they are the current leaf (so the active position stays visible).
+			"user": true, "assistant": true, "tool": false, "leafTool": true,
+			"modelChange": false, "label": false, "sessionInfo": false,
+			"compaction": true, "labelled": true,
+		},
+		TreeFilterNoTools: {
+			"user": true, "assistant": true, "tool": false, "leafTool": false,
+			"modelChange": true, "label": true, "sessionInfo": true,
+			"compaction": true, "labelled": true,
+		},
+		TreeFilterUserOnly: {
+			"user": true, "assistant": false, "tool": false, "leafTool": false,
+			"modelChange": false, "label": false, "sessionInfo": false,
+			"compaction": false, "labelled": false,
+		},
+		TreeFilterLabelOnly: {
+			"user": false, "assistant": false, "tool": false, "leafTool": false,
+			"modelChange": false, "label": false, "sessionInfo": false,
+			"compaction": false, "labelled": true,
+		},
+	}
+
+	for filter, expectations := range want {
+		ts := &TreeSelectorComponent{filter: filter, leafID: leafID}
+		for key, visible := range expectations {
+			if got := ts.passesFilter(nodes[key]); got != visible {
+				t.Errorf("filter %s: passesFilter(%s) = %v, want %v", filter, key, got, visible)
+			}
+		}
+	}
+}
+
+func TestTreeSelectorFlattensNestedTree(t *testing.T) {
+	// user → assistant → tool(leaf), plus a sibling branch off the assistant.
+	tree := []app.TreeNodeView{{
+		ID: "u1", Kind: app.EntryKindMessage, Role: "user", Text: "hi",
+		Children: []app.TreeNodeView{{
+			ID: "a1", ParentID: "u1", Kind: app.EntryKindMessage, Role: "assistant", Text: "hello",
+			Children: []app.TreeNodeView{
+				{ID: "t1", ParentID: "a1", Kind: app.EntryKindMessage, Role: "tool"},
+				{ID: "u2", ParentID: "a1", Kind: app.EntryKindMessage, Role: "user", Text: "more"},
+			},
+		}},
+	}}
+
+	ts := NewTreeSelector(tree, "u2", 80, 24)
+
+	var ids []string
+	for _, n := range ts.flatNodes {
+		ids = append(ids, n.ID)
+	}
+	// Default filter hides the non-leaf tool message.
+	if got := strings.Join(ids, ","); got != "u1,a1,u2" {
+		t.Errorf("flattened ids = %q, want %q", got, "u1,a1,u2")
+	}
+
+	// Depth drives the indentation prefix; roots are never indented.
+	if ts.flatNodes[0].Depth != 0 || ts.flatNodes[0].Prefix != "" {
+		t.Errorf("root node depth/prefix = %d/%q, want 0/\"\"", ts.flatNodes[0].Depth, ts.flatNodes[0].Prefix)
+	}
+	if ts.flatNodes[1].Depth != 1 {
+		t.Errorf("child depth = %d, want 1", ts.flatNodes[1].Depth)
+	}
+	if ts.flatNodes[2].Depth != 2 {
+		t.Errorf("grandchild depth = %d, want 2", ts.flatNodes[2].Depth)
+	}
+
+	// The cursor starts on the active leaf.
+	if got := ts.flatNodes[ts.popup.Cursor()].ID; got != "u2" {
+		t.Errorf("cursor is on %q, want the leaf %q", got, "u2")
+	}
+}
+
+func TestNewTreeSelectorForForkStartsOnLastUserMessage(t *testing.T) {
+	tree := []app.TreeNodeView{{
+		ID: "u1", Kind: app.EntryKindMessage, Role: "user", Text: "first",
+		Children: []app.TreeNodeView{{
+			ID: "a1", ParentID: "u1", Kind: app.EntryKindMessage, Role: "assistant", Text: "reply",
+			Children: []app.TreeNodeView{{
+				ID: "u2", ParentID: "a1", Kind: app.EntryKindMessage, Role: "user", Text: "second",
+			}},
+		}},
+	}}
+
+	ts := NewTreeSelectorForFork(tree, "u2", 80, 24)
+
+	// Fork mode lists user messages only.
+	if len(ts.flatNodes) != 2 {
+		t.Fatalf("len(flatNodes) = %d, want 2: %+v", len(ts.flatNodes), ts.flatNodes)
+	}
+	selected := ts.flatNodes[ts.popup.Cursor()]
+	if selected.ID != "u2" {
+		t.Errorf("cursor is on %q, want the last user message %q", selected.ID, "u2")
+	}
+	// ParentID is what a fork branches from, so it must survive flattening.
+	if selected.ParentID != "a1" {
+		t.Errorf("ParentID = %q, want %q", selected.ParentID, "a1")
+	}
+}

From 61a94ebc9a5ed966891292107701c0b962de9dd8 Mon Sep 17 00:00:00 2001
From: Ed Zynda <ezynda3@gmail.com>
Date: Thu, 30 Jul 2026 16:34:57 +0300
Subject: [PATCH 2/2] fix(ui): address CodeRabbit review on session snapshot
 refactor (#101)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- renderSessionHistory no longer returns early on an empty history, which
  left the previous conversation on screen. The old code returned when the
  branch had no *entries*; the refactor returned when it had no *messages*,
  so a branch holding only non-message entries stopped clearing. This is
  reachable from /retry and /undo (both pop the last user message and can
  empty the branch), and from forking or resuming into an empty session.
  The no-session guard is preserved via SessionSnapshot so an in-memory
  run's transcript is still left untouched.

- Wrap propagated errors in SetSessionName, NewSession, ForkSession and
  SessionSystemPromptEntry with fmt.Errorf/%w so failures are diagnosable
  across the app/UI boundary.

Skipped: adding context.Context to NewSession/ForkSession/
SessionSystemPromptEntry. Nothing beneath them accepts a context —
internal/session performs synchronous file I/O with no cancellation — and
no method on AppController takes one, so the parameter would be ignored
and would advertise cancellation that does not exist. Threading context
through the session layer is its own change.

Refs #101
---
 internal/app/session_view.go               |  17 +++-
 internal/ui/model.go                       |  12 ++-
 internal/ui/model_test.go                  |  12 ++-
 internal/ui/render_session_history_test.go | 111 +++++++++++++++++++++
 4 files changed, 143 insertions(+), 9 deletions(-)
 create mode 100644 internal/ui/render_session_history_test.go

diff --git a/internal/app/session_view.go b/internal/app/session_view.go
index b340dc74..968eb123 100644
--- a/internal/app/session_view.go
+++ b/internal/app/session_view.go
@@ -2,6 +2,7 @@ package app
 
 import (
 	"errors"
+	"fmt"
 	"time"
 
 	"github.com/mark3labs/kit/internal/message"
@@ -225,8 +226,10 @@ func (a *App) SetSessionName(name string) error {
 	if tm == nil {
 		return ErrNoSession
 	}
-	_, err := tm.AppendSessionInfo(name)
-	return err
+	if _, err := tm.AppendSessionInfo(name); err != nil {
+		return fmt.Errorf("append session info: %w", err)
+	}
+	return nil
 }
 
 // NewSession replaces the active tree session with a brand new one created in
@@ -242,7 +245,7 @@ func (a *App) NewSession(cwd string) error {
 	}
 	ts, err := session.CreateTreeSession(cwd)
 	if err != nil {
-		return err
+		return fmt.Errorf("create tree session: %w", err)
 	}
 	a.SwitchTreeSession(ts)
 	return nil
@@ -258,7 +261,7 @@ func (a *App) ForkSession(cwd, targetID string) error {
 	}
 	ts, err := tm.ForkToNewSession(cwd, targetID)
 	if err != nil {
-		return err
+		return fmt.Errorf("fork session from entry %q: %w", targetID, err)
 	}
 	a.SwitchTreeSession(ts)
 	return nil
@@ -280,5 +283,9 @@ func (a *App) SessionSystemPromptEntry(systemPrompt, fallbackModelID string) ([]
 	if modelID == "" {
 		modelID = fallbackModelID
 	}
-	return session.MarshalEntry(session.NewSystemPromptEntry(systemPrompt, modelID, provider))
+	data, err := session.MarshalEntry(session.NewSystemPromptEntry(systemPrompt, modelID, provider))
+	if err != nil {
+		return nil, fmt.Errorf("marshal system prompt entry: %w", err)
+	}
+	return data, nil
 }
diff --git a/internal/ui/model.go b/internal/ui/model.go
index 139eb60a..04c49dc5 100644
--- a/internal/ui/model.go
+++ b/internal/ui/model.go
@@ -5811,12 +5811,20 @@ func (m *AppModel) handleResumeCommand() tea.Cmd {
 // This gives the user visual context of the conversation when resuming or
 // importing a session. Call this after switchSession succeeds.
 func (m *AppModel) renderSessionHistory() {
-	history := m.appCtrl.SessionHistory()
-	if len(history) == 0 {
+	if _, ok := m.appCtrl.SessionSnapshot(); !ok {
+		// No session to render from — leave the transcript untouched rather
+		// than blanking a conversation the session layer knows nothing about
+		// (e.g. an in-memory run).
 		return
 	}
 
+	history := m.appCtrl.SessionHistory()
+
 	// Clear existing messages so we start fresh with the resumed session.
+	// This must happen even when the history is empty: /retry and /undo pop
+	// the last user message and can leave the branch with no messages at all,
+	// and forking or resuming can land on an empty session. Returning early
+	// here would leave the previous transcript on screen.
 	m.messages = []MessageItem{}
 
 	// First pass: build a map of tool call ID → {name, args} from assistant
diff --git a/internal/ui/model_test.go b/internal/ui/model_test.go
index 315d09b0..63e7a58c 100644
--- a/internal/ui/model_test.go
+++ b/internal/ui/model_test.go
@@ -25,6 +25,11 @@ type stubAppController struct {
 	clearQueueCalled int
 	clearMsgCalled   int
 	queueLen         int
+
+	// hasSession controls whether SessionSnapshot reports an active session,
+	// and sessionHistory is what SessionHistory returns for it.
+	hasSession     bool
+	sessionHistory []message.Message
 }
 
 func (s *stubAppController) Run(prompt string) int {
@@ -58,7 +63,10 @@ func (s *stubAppController) CompactConversation(_ string) error {
 }
 
 func (s *stubAppController) SessionSnapshot() (app.SessionSnapshot, bool) {
-	return app.SessionSnapshot{}, false
+	if !s.hasSession {
+		return app.SessionSnapshot{}, false
+	}
+	return app.SessionSnapshot{ID: "stub-session"}, true
 }
 
 func (s *stubAppController) SessionTree() []app.TreeNodeView {
@@ -66,7 +74,7 @@ func (s *stubAppController) SessionTree() []app.TreeNodeView {
 }
 
 func (s *stubAppController) SessionHistory() []message.Message {
-	return nil
+	return s.sessionHistory
 }
 
 func (s *stubAppController) SetSessionName(_ string) error {
diff --git a/internal/ui/render_session_history_test.go b/internal/ui/render_session_history_test.go
new file mode 100644
index 00000000..bf965de9
--- /dev/null
+++ b/internal/ui/render_session_history_test.go
@@ -0,0 +1,111 @@
+package ui
+
+import (
+	"testing"
+
+	"github.com/mark3labs/kit/internal/message"
+)
+
+// textMessage builds a single-text-part message for history fixtures.
+func textMessage(role message.MessageRole, text string) message.Message {
+	return message.Message{
+		Role:  role,
+		Parts: []message.ContentPart{message.TextContent{Text: text}},
+	}
+}
+
+// seedTranscript puts a stale message into the model's visible transcript so
+// tests can assert whether renderSessionHistory clears it.
+func seedTranscript(m *AppModel) {
+	m.messages = []MessageItem{
+		NewThemedMessageItem(generateMessageID(), "user", "stale message", func() string {
+			return "stale message"
+		}),
+	}
+}
+
+// An empty branch must still clear the transcript. /retry and /undo pop the
+// last user message and can leave the branch with no messages at all; forking
+// or resuming can land on an empty session. Returning early there left the
+// previous conversation on screen.
+func TestRenderSessionHistoryClearsOnEmptyHistory(t *testing.T) {
+	ctrl := &stubAppController{hasSession: true, sessionHistory: nil}
+	m, _, _ := newTestAppModel(ctrl)
+	seedTranscript(m)
+
+	m.renderSessionHistory()
+
+	if len(m.messages) != 0 {
+		t.Fatalf("stale transcript survived an empty history: %d messages remain", len(m.messages))
+	}
+	if !m.layoutDirty {
+		t.Error("layoutDirty = false, want true so the cleared list is re-laid out")
+	}
+	if !m.pendingGotoBottom {
+		t.Error("pendingGotoBottom = false, want true")
+	}
+}
+
+// Without a session there is nothing to render from, so the transcript must be
+// left alone rather than blanked.
+func TestRenderSessionHistoryLeavesTranscriptWhenNoSession(t *testing.T) {
+	ctrl := &stubAppController{hasSession: false}
+	m, _, _ := newTestAppModel(ctrl)
+	seedTranscript(m)
+
+	m.renderSessionHistory()
+
+	if len(m.messages) != 1 {
+		t.Fatalf("transcript was modified without an active session: %d messages", len(m.messages))
+	}
+}
+
+func TestRenderSessionHistoryRendersMessages(t *testing.T) {
+	ctrl := &stubAppController{
+		hasSession: true,
+		sessionHistory: []message.Message{
+			textMessage(message.RoleUser, "what is the capital of France?"),
+			textMessage(message.RoleAssistant, "Paris."),
+		},
+	}
+	m, _, _ := newTestAppModel(ctrl)
+	seedTranscript(m)
+
+	m.renderSessionHistory()
+
+	if len(m.messages) != 2 {
+		t.Fatalf("len(messages) = %d, want 2", len(m.messages))
+	}
+	// The stale entry must be gone, replaced by the session's own history.
+	first, ok := m.messages[0].(*TextMessageItem)
+	if !ok {
+		t.Fatalf("messages[0] is %T, want *TextMessageItem", m.messages[0])
+	}
+	if first.role != "user" || first.content != "what is the capital of France?" {
+		t.Errorf("messages[0] = %s/%q, want user/the seeded question", first.role, first.content)
+	}
+	second, ok := m.messages[1].(*TextMessageItem)
+	if !ok {
+		t.Fatalf("messages[1] is %T, want *TextMessageItem", m.messages[1])
+	}
+	if second.role != "assistant" || second.content != "Paris." {
+		t.Errorf("messages[1] = %s/%q, want assistant/Paris.", second.role, second.content)
+	}
+}
+
+// Messages with no text content (pure tool interactions) contribute no
+// transcript rows, but must still clear whatever was there before.
+func TestRenderSessionHistoryClearsWhenHistoryHasNoRenderableText(t *testing.T) {
+	ctrl := &stubAppController{
+		hasSession:     true,
+		sessionHistory: []message.Message{textMessage(message.RoleUser, "   ")},
+	}
+	m, _, _ := newTestAppModel(ctrl)
+	seedTranscript(m)
+
+	m.renderSessionHistory()
+
+	if len(m.messages) != 0 {
+		t.Fatalf("expected an empty transcript, got %d messages", len(m.messages))
+	}
+}