1package modes
2
3import (
4 "fmt"
5 "strings"
6
7 "github.com/patriceckhart/zot/packages/core"
8 "github.com/patriceckhart/zot/packages/provider"
9 "github.com/patriceckhart/zot/packages/tui"
10)
11
12// sessionTreeDialog renders a compact outline of the current session family.
13// Branches are shown inline at the point where they forked, using indentation
14// rather than file-level rows. Selecting a node checks out that point into a
15// new branch.
16type sessionTreeDialog struct {
17 active bool
18 items []sessionTreeItem
19 cursor int
20}
21
22type sessionTreeItem struct {
23 label string
24 messageIdx int
25 turnNo int
26 role provider.Role
27 prompt string
28 path string
29 depth int
30 current bool
31}
32
33type sessionTreeAction struct {
34 Select bool
35 MessageIdx int
36 TurnNo int
37 Role provider.Role
38 Prompt string
39 Path string
40 Close bool
41}
42
43func newSessionTreeDialog() *sessionTreeDialog { return &sessionTreeDialog{} }
44
45func (d *sessionTreeDialog) OpenMessages(msgs []provider.Message) bool {
46 d.items = buildSessionTreeItems("", msgs, 0, false)
47 d.cursor = len(d.items) - 1
48 if d.cursor < 0 {
49 d.cursor = 0
50 }
51 d.active = true
52 return len(d.items) > 0
53}
54
55func (d *sessionTreeDialog) OpenSessionFamily(root, cwd, currentPath string) bool {
56 roots := core.BuildSessionTree(root, cwd)
57 if len(roots) == 0 {
58 return false
59 }
60 rootNode := findTreeRootForPath(roots, currentPath)
61 if rootNode == nil {
62 rootNode = roots[0]
63 }
64 d.items = flattenSessionFamily(rootNode, currentPath)
65 d.cursor = indexCurrentTreeItem(d.items)
66 d.active = true
67 return len(d.items) > 0
68}
69
70// Close hides the dialog.
71func (d *sessionTreeDialog) Close() { d.active = false }
72
73// CursorPos returns -1 because this dialog has no inline editor.
74func (d *sessionTreeDialog) CursorPos() (row, col int) { return -1, -1 }
75
76// Active reports whether the dialog consumes input.
77func (d *sessionTreeDialog) Active() bool { return d != nil && d.active }
78
79// HandleKey advances the cursor or resolves the selection.
80func (d *sessionTreeDialog) HandleKey(k tui.Key) sessionTreeAction {
81 switch k.Kind {
82 case tui.KeyUp:
83 if d.cursor > 0 {
84 d.cursor--
85 }
86 case tui.KeyDown:
87 if d.cursor < len(d.items)-1 {
88 d.cursor++
89 }
90 case tui.KeyPageUp:
91 d.cursor -= 5
92 if d.cursor < 0 {
93 d.cursor = 0
94 }
95 case tui.KeyPageDown:
96 d.cursor += 5
97 if d.cursor >= len(d.items) {
98 d.cursor = len(d.items) - 1
99 }
100 if d.cursor < 0 {
101 d.cursor = 0
102 }
103 case tui.KeyEsc:
104 d.Close()
105 return sessionTreeAction{Close: true}
106 case tui.KeyEnter:
107 if len(d.items) == 0 || d.cursor < 0 || d.cursor >= len(d.items) {
108 d.Close()
109 return sessionTreeAction{Close: true}
110 }
111 it := d.items[d.cursor]
112 d.Close()
113 return sessionTreeAction{Select: true, MessageIdx: it.messageIdx, TurnNo: it.turnNo, Role: it.role, Prompt: it.prompt, Path: it.path}
114 }
115 return sessionTreeAction{}
116}
117
118// Render returns the dialog lines.
119func (d *sessionTreeDialog) Render(th tui.Theme, width int) []string {
120 if !d.Active() {
121 return nil
122 }
123 var lines []string
124 lines = append(lines, frameHeader(th, "session tree", width))
125 if len(d.items) == 0 {
126 lines = append(lines, th.FG256(th.Muted, "no messages in this session yet"))
127 lines = append(lines, th.FG256(th.Muted, "press esc to close"))
128 lines = append(lines, frameRule(th, width))
129 return lines
130 }
131 lines = append(lines, th.FG256(th.Muted, "session branches (↑/↓, pgup/pgdn, enter checkout, esc cancel):"))
132
133 const maxRows = 12
134 start := 0
135 end := len(d.items)
136 if end > maxRows {
137 start = d.cursor - maxRows/2
138 if start < 0 {
139 start = 0
140 }
141 end = start + maxRows
142 if end > len(d.items) {
143 end = len(d.items)
144 start = end - maxRows
145 }
146 }
147 if start > 0 {
148 lines = append(lines, th.FG256(th.Muted, fmt.Sprintf(" ↑ %d more above", start)))
149 }
150 for i := start; i < end; i++ {
151 it := d.items[i]
152 indent := strings.Repeat(" ", it.depth)
153 plain := " " + indent + fitSessionTreeLabel(it.label, width-2-len([]rune(indent)))
154 if it.current {
155 plain += " [current]"
156 }
157 if i == d.cursor {
158 lines = append(lines, th.PadHighlight(plain, width))
159 } else {
160 lines = append(lines, colorSessionTreeLine(th, plain))
161 }
162 }
163 if end < len(d.items) {
164 lines = append(lines, th.FG256(th.Muted, fmt.Sprintf(" ↓ %d more below", len(d.items)-end)))
165 }
166 lines = append(lines, th.FG256(th.Muted, fmt.Sprintf("%d/%d", d.cursor+1, len(d.items))))
167 lines = append(lines, frameRule(th, width))
168 return lines
169}
170
171func flattenSessionFamily(root *core.TreeNode, currentPath string) []sessionTreeItem {
172 childrenByFork := map[int][]*core.TreeNode{}
173 for _, child := range root.Children {
174 childrenByFork[child.Meta.ForkPoint] = append(childrenByFork[child.Meta.ForkPoint], child)
175 }
176 sess, msgs, err := core.OpenSession(root.Summary.Path)
177 if err != nil {
178 return nil
179 }
180 _ = sess.Close()
181 items := buildSessionTreeItems(root.Summary.Path, msgs, 0, root.Summary.Path == currentPath)
182 var out []sessionTreeItem
183 for idx, item := range items {
184 out = append(out, item)
185 for _, child := range childrenByFork[idx+1] {
186 out = append(out, flattenSessionBranch(child, currentPath, 1)...)
187 }
188 }
189 for _, child := range childrenByFork[0] {
190 out = append(out, flattenSessionBranch(child, currentPath, 1)...)
191 }
192 return out
193}
194
195func flattenSessionBranch(node *core.TreeNode, currentPath string, depth int) []sessionTreeItem {
196 sess, msgs, err := core.OpenSession(node.Summary.Path)
197 if err != nil {
198 return nil
199 }
200 _ = sess.Close()
201 start := node.Meta.ForkPoint
202 if start < 0 {
203 start = 0
204 }
205 if start > len(msgs) {
206 start = len(msgs)
207 }
208 items := buildSessionTreeItems(node.Summary.Path, msgs[start:], depth, node.Summary.Path == currentPath)
209 for idx := range items {
210 items[idx].messageIdx += start
211 }
212 childrenByFork := map[int][]*core.TreeNode{}
213 for _, child := range node.Children {
214 childrenByFork[child.Meta.ForkPoint] = append(childrenByFork[child.Meta.ForkPoint], child)
215 }
216 var out []sessionTreeItem
217 for relIdx, item := range items {
218 out = append(out, item)
219 absAfter := start + relIdx + 1
220 for _, child := range childrenByFork[absAfter] {
221 out = append(out, flattenSessionBranch(child, currentPath, depth+1)...)
222 }
223 }
224 return out
225}
226
227func buildSessionTreeItems(path string, msgs []provider.Message, depth int, currentPath bool) []sessionTreeItem {
228 out := make([]sessionTreeItem, 0, len(msgs))
229 turn := 0
230 lastTurn := 0
231 for idx, msg := range msgs {
232 label := sessionTreeRoleLabel(msg.Role)
233 if msg.Role == provider.RoleUser {
234 turn++
235 lastTurn = turn
236 } else if lastTurn == 0 {
237 lastTurn = 1
238 }
239 preview := sessionTreePreview(msg)
240 out = append(out, sessionTreeItem{
241 label: fmt.Sprintf("%s: %s", label, preview),
242 messageIdx: idx,
243 turnNo: lastTurn,
244 role: msg.Role,
245 prompt: firstTextFromTreeMessage(msg),
246 path: path,
247 depth: depth,
248 current: currentPath && idx == len(msgs)-1,
249 })
250 }
251 return out
252}
253
254func findTreeRootForPath(roots []*core.TreeNode, path string) *core.TreeNode {
255 for _, root := range roots {
256 if treeContainsPath(root, path) {
257 return root
258 }
259 }
260 return nil
261}
262
263func treeContainsPath(n *core.TreeNode, path string) bool {
264 if n == nil {
265 return false
266 }
267 if n.Summary.Path == path {
268 return true
269 }
270 for _, child := range n.Children {
271 if treeContainsPath(child, path) {
272 return true
273 }
274 }
275 return false
276}
277
278func indexCurrentTreeItem(items []sessionTreeItem) int {
279 for i, it := range items {
280 if it.current {
281 return i
282 }
283 }
284 if len(items) == 0 {
285 return 0
286 }
287 return len(items) - 1
288}
289
290func firstTextFromTreeMessage(msg provider.Message) string {
291 for _, c := range msg.Content {
292 if tb, ok := c.(provider.TextBlock); ok {
293 return tb.Text
294 }
295 }
296 return ""
297}
298
299func sessionTreeRoleLabel(role provider.Role) string {
300 switch role {
301 case provider.RoleUser:
302 return "you"
303 case provider.RoleAssistant:
304 return "zot"
305 case provider.RoleTool:
306 return "tool"
307 default:
308 return string(role)
309 }
310}
311
312func sessionTreePreview(msg provider.Message) string {
313 var parts []string
314 for _, c := range msg.Content {
315 switch b := c.(type) {
316 case provider.TextBlock:
317 text := strings.TrimSpace(strings.ReplaceAll(b.Text, "\n", " "))
318 if text != "" {
319 parts = append(parts, text)
320 }
321 case provider.ImageBlock:
322 parts = append(parts, "[image]")
323 case provider.ToolCallBlock:
324 parts = append(parts, "tool "+b.Name)
325 case provider.ToolResultBlock:
326 if b.IsError {
327 parts = append(parts, "tool result error")
328 } else {
329 parts = append(parts, "tool result")
330 }
331 case provider.ReasoningBlock:
332 if b.Summary != "" {
333 parts = append(parts, "reasoning")
334 }
335 }
336 }
337 if len(parts) == 0 {
338 return "(empty)"
339 }
340 return strings.Join(parts, " ")
341}
342
343func colorSessionTreeLine(th tui.Theme, line string) string {
344 plain := strings.TrimSpace(line)
345 switch {
346 case strings.HasPrefix(plain, "you:"):
347 return th.FG256(th.FG, line)
348 case strings.HasPrefix(plain, "zot:"):
349 return th.FG256(th.Muted, line)
350 case strings.HasPrefix(plain, "tool:"):
351 return th.FG256(th.ToolOut, line)
352 default:
353 return th.FG256(th.Muted, line)
354 }
355}
356
357func fitSessionTreeLabel(label string, maxWidth int) string {
358 if maxWidth < 4 {
359 maxWidth = 4
360 }
361 runes := []rune(label)
362 if len(runes) <= maxWidth {
363 return label
364 }
365 return string(runes[:maxWidth-3]) + "..."
366}