1package modes
2
3import (
4 "os"
5 "strings"
6 "testing"
7 "time"
8
9 "github.com/bnema/zut/packages/core"
10 "github.com/bnema/zut/packages/provider"
11 "github.com/bnema/zut/packages/tui"
12)
13
14func TestSessionTreeTargetCarriesEffectiveBoundaryAndDraft(t *testing.T) {
15 msgs := []provider.Message{
16 treeTestMessage(provider.RoleUser, provider.TextBlock{Text: "first"}),
17 treeTestMessage(provider.RoleAssistant, provider.TextBlock{Text: "answer"}),
18 treeTestMessage(provider.RoleUser,
19 provider.TextBlock{Text: "line one"},
20 provider.ImageBlock{MimeType: "image/png", Data: []byte{1, 2}},
21 provider.TextBlock{Text: "line two"}),
22 }
23
24 d := newSessionTreeDialog()
25 if !d.OpenMessages(msgs) {
26 t.Fatal("OpenMessages returned false")
27 }
28 d.cursor = 2
29 act := d.HandleKey(tui.Key{Kind: tui.KeyEnter})
30 if !act.Select {
31 t.Fatal("enter did not select a target")
32 }
33 got := act.Target
34 if got.SourcePath != "" || got.EffectiveIndex != 2 || got.SelectionBoundary != 2 {
35 t.Fatalf("target indices/path = %+v", got)
36 }
37 if got.Role != provider.RoleUser || got.Boundary != sessionTreeMessageBoundary {
38 t.Fatalf("target role/boundary = %+v", got)
39 }
40 if got.UserDraft != "line one\nline two" {
41 t.Fatalf("user draft = %q", got.UserDraft)
42 }
43 // The scalar fields remain usable by the current integration layer.
44 if act.MessageIdx != 2 || act.Role != provider.RoleUser || act.Prompt != got.UserDraft {
45 t.Fatalf("legacy action fields = %+v", act)
46 }
47}
48
49func TestSessionTreeRetainsPreCompactionMessagesAndHidesToolRows(t *testing.T) {
50 root := t.TempDir()
51 cwd := "/workspace/tree-history"
52 session, err := core.NewSession(root, cwd, "test", "model", "test")
53 if err != nil {
54 t.Fatal(err)
55 }
56 for _, message := range []provider.Message{
57 treeTestMessage(provider.RoleUser, provider.TextBlock{Text: "old prompt"}),
58 treeTestMessage(provider.RoleAssistant, provider.TextBlock{Text: "old answer"}),
59 {Role: provider.RoleAssistant, Content: []provider.Content{provider.ToolCallBlock{ID: "call-1", Name: "bash"}}},
60 {Role: provider.RoleTool, Content: []provider.Content{provider.ToolResultBlock{CallID: "call-1", Content: []provider.Content{provider.TextBlock{Text: "tool output"}}}}},
61 treeTestMessage(provider.RoleAssistant, provider.TextBlock{Text: "old final"}),
62 } {
63 if err := session.AppendMessage(message); err != nil {
64 t.Fatal(err)
65 }
66 }
67 if err := session.AppendCompaction([]provider.Message{
68 {
69 Role: provider.RoleUser,
70 Meta: map[string]string{"compaction": "true"},
71 Content: []provider.Content{provider.TextBlock{Text: "summary"}},
72 },
73 treeTestMessage(provider.RoleAssistant, provider.TextBlock{Text: "old final"}),
74 }); err != nil {
75 t.Fatal(err)
76 }
77 if err := session.AppendMessage(treeTestMessage(provider.RoleAssistant, provider.TextBlock{Text: "new answer"})); err != nil {
78 t.Fatal(err)
79 }
80 if err := session.AppendMessage(provider.Message{
81 Role: provider.RoleUser,
82 Meta: map[string]string{shellEscapeMetaKey: "true"},
83 Content: []provider.Content{provider.TextBlock{Text: "shell output"}},
84 }); err != nil {
85 t.Fatal(err)
86 }
87 path := session.Path
88 if err := session.Close(); err != nil {
89 t.Fatal(err)
90 }
91
92 d := newSessionTreeDialog()
93 if !d.OpenSessionFamily(root, cwd, path) {
94 t.Fatal("OpenSessionFamily returned false")
95 }
96 var oldPrompt, newAnswer *sessionTreeItem
97 oldFinalRows := 0
98 for idx := range d.items {
99 item := &d.items[idx]
100 if strings.Contains(item.label, "old prompt") {
101 oldPrompt = item
102 }
103 if strings.Contains(item.label, "new answer") {
104 newAnswer = item
105 }
106 if strings.Contains(item.label, "old final") {
107 oldFinalRows++
108 }
109 if strings.Contains(item.label, "tool") || strings.Contains(item.label, "result") || strings.Contains(item.label, "summary") || strings.Contains(item.label, "shell output") {
110 t.Fatalf("internal/tool row leaked into fork picker: %q", item.label)
111 }
112 }
113 if oldPrompt == nil || !oldPrompt.target.Historical {
114 t.Fatalf("old pre-compaction prompt target = %+v, want historical target", oldPrompt)
115 }
116 if newAnswer == nil || newAnswer.target.Historical {
117 t.Fatalf("current answer target = %+v, want effective target", newAnswer)
118 }
119 if oldFinalRows != 1 {
120 t.Fatalf("preserved compaction-tail row rendered %d times, want once", oldFinalRows)
121 }
122}
123
124func TestSessionTreeHistoricalSelectionCreatesBranchFromOldSegment(t *testing.T) {
125 root := t.TempDir()
126 cwd := "/workspace/tree-history-branch"
127 session, err := core.NewSession(root, cwd, "test", "model", "test")
128 if err != nil {
129 t.Fatal(err)
130 }
131 if err := session.AppendMessage(treeTestMessage(provider.RoleUser, provider.TextBlock{Text: "old prompt"})); err != nil {
132 t.Fatal(err)
133 }
134 if err := session.AppendMessage(treeTestMessage(provider.RoleAssistant, provider.TextBlock{Text: "old answer"})); err != nil {
135 t.Fatal(err)
136 }
137 if err := session.AppendCompaction([]provider.Message{{
138 Role: provider.RoleUser,
139 Meta: map[string]string{"compaction": "true"},
140 Content: []provider.Content{provider.TextBlock{Text: "summary"}},
141 }}); err != nil {
142 t.Fatal(err)
143 }
144 if err := session.AppendMessage(treeTestMessage(provider.RoleAssistant, provider.TextBlock{Text: "new answer"})); err != nil {
145 t.Fatal(err)
146 }
147 path := session.Path
148 if err := session.Close(); err != nil {
149 t.Fatal(err)
150 }
151
152 d := newSessionTreeDialog()
153 if !d.OpenSessionFamily(root, cwd, path) {
154 t.Fatal("OpenSessionFamily returned false")
155 }
156 var target sessionTreeTarget
157 found := false
158 for _, item := range d.items {
159 if strings.Contains(item.label, "old answer") {
160 target = item.target
161 found = true
162 break
163 }
164 }
165 if !found || !target.Historical {
166 t.Fatalf("old target = %+v, found=%v", target, found)
167 }
168
169 ag := core.NewAgent(nil, "model", "", nil)
170 var loaded string
171 i := NewInteractive(InteractiveConfig{
172 Agent: ag,
173 CWD: cwd,
174 SessionsRoot: root,
175 Version: "test",
176 CurrentSessionPath: func() string { return path },
177 LoadSession: func(newPath string) error {
178 loaded = newPath
179 return nil
180 },
181 })
182 i.applySessionTreeTarget(target, 1)
183 if loaded == "" {
184 t.Fatalf("historical selection did not load a branch; status=%q", i.statusErr)
185 }
186 branch, messages, err := core.OpenSession(loaded)
187 if err != nil {
188 t.Fatal(err)
189 }
190 if err := branch.Close(); err != nil {
191 t.Fatal(err)
192 }
193 if len(messages) != 2 || userMessageText(messages[0]) != "old prompt" || userMessageText(messages[1]) != "old answer" {
194 t.Fatalf("historical branch messages = %+v, want old prompt and answer", messages)
195 }
196}
197
198func TestSessionTreeUsesCurrentFamilyAndNoRootFallback(t *testing.T) {
199 root := t.TempDir()
200 cwd := "/workspace/current-family"
201 first := newTreeTestSession(t, root, cwd, []provider.Message{
202 treeTestMessage(provider.RoleUser, provider.TextBlock{Text: "unrelated root"}),
203 })
204 unrelatedFile, err := os.OpenFile(first, os.O_APPEND|os.O_WRONLY, 0o644)
205 if err != nil {
206 t.Fatal(err)
207 }
208 if _, err := unrelatedFile.WriteString("{unrelated corruption}\n"); err != nil {
209 _ = unrelatedFile.Close()
210 t.Fatal(err)
211 }
212 if err := unrelatedFile.Close(); err != nil {
213 t.Fatal(err)
214 }
215 current := newTreeTestSession(t, root, cwd, []provider.Message{
216 treeTestMessage(provider.RoleUser, provider.TextBlock{Text: "current root"}),
217 })
218
219 d := newSessionTreeDialog()
220 if !d.OpenSessionFamily(root, cwd, current) {
221 t.Fatal("OpenSessionFamily returned false")
222 }
223 for _, item := range d.items {
224 if item.path != current {
225 t.Fatalf("item from unrelated family: %+v (first root %q)", item, first)
226 }
227 }
228 if !d.Active() {
229 t.Fatal("dialog is not active")
230 }
231 oldItems := append([]sessionTreeItem(nil), d.items...)
232 if d.OpenSessionFamily(root, cwd, t.TempDir()+"/not-a-session.jsonl") {
233 t.Fatal("missing current path selected a forest root")
234 }
235 if !d.Active() || len(d.items) != len(oldItems) || d.items[0].path != oldItems[0].path {
236 t.Fatal("failed open changed the active dialog")
237 }
238}
239
240func TestSessionTreeShowsEmptyAndDetachedBoundaries(t *testing.T) {
241 root := t.TempDir()
242 cwd := "/workspace/stale-compaction"
243 parent := newTreeTestSession(t, root, cwd, []provider.Message{
244 treeTestMessage(provider.RoleUser, provider.TextBlock{Text: "one"}),
245 treeTestMessage(provider.RoleAssistant, provider.TextBlock{Text: "two"}),
246 treeTestMessage(provider.RoleUser, provider.TextBlock{Text: "three"}),
247 })
248
249 // Both children are created against the old effective parent. The first
250 // has no post-fork rows; the second's fork point becomes stale after the
251 // parent is compacted to one message.
252 emptyChild, err := core.BranchSessionHidden(parent, root, cwd, "test", 1)
253 if err != nil {
254 t.Fatal(err)
255 }
256 staleChild, err := core.BranchSessionHidden(parent, root, cwd, "test", 3)
257 if err != nil {
258 t.Fatal(err)
259 }
260 parentSession, _, err := core.OpenSession(parent)
261 if err != nil {
262 t.Fatal(err)
263 }
264 if err := parentSession.AppendCompaction([]provider.Message{
265 treeTestMessage(provider.RoleUser, provider.TextBlock{Text: "compacted"}),
266 }); err != nil {
267 t.Fatal(err)
268 }
269 if err := parentSession.Close(); err != nil {
270 t.Fatal(err)
271 }
272
273 d := newSessionTreeDialog()
274 if !d.OpenSessionFamily(root, cwd, parent) {
275 t.Fatal("OpenSessionFamily returned false after compaction")
276 }
277 var empty, detached *sessionTreeItem
278 staleMessageRows := 0
279 for i := range d.items {
280 item := &d.items[i]
281 switch {
282 case item.target.IsEmptyBoundary() && item.path == emptyChild:
283 empty = item
284 case item.target.IsDetachedBoundary() && item.path == staleChild:
285 detached = item
286 case item.path == staleChild && !item.target.IsBoundary():
287 staleMessageRows++
288 }
289 }
290 if empty == nil || detached == nil {
291 t.Fatalf("boundaries missing from tree: empty=%v detached=%v items=%+v", empty != nil, detached != nil, d.items)
292 }
293 if staleMessageRows != 3 {
294 t.Fatalf("detached branch rendered %d message rows, want its complete snapshot (3)", staleMessageRows)
295 }
296 if empty.target.SelectionBoundary != 1 || detached.target.SelectionBoundary != 3 {
297 t.Fatalf("boundary indices = empty %+v detached %+v", empty.target, detached.target)
298 }
299 if !strings.Contains(strings.ToLower(empty.label), "empty") || !strings.Contains(strings.ToLower(detached.label), "detached") {
300 t.Fatalf("boundary labels = %q / %q", empty.label, detached.label)
301 }
302}
303
304func TestSessionTreePreflightIsAtomicOnMalformedFamilyMember(t *testing.T) {
305 root := t.TempDir()
306 cwd := "/workspace/atomic-tree"
307 parent := newTreeTestSession(t, root, cwd, []provider.Message{
308 treeTestMessage(provider.RoleUser, provider.TextBlock{Text: "parent"}),
309 })
310 child, err := core.BranchSessionHidden(parent, root, cwd, "test", 1)
311 if err != nil {
312 t.Fatal(err)
313 }
314 f, err := os.OpenFile(child, os.O_APPEND|os.O_WRONLY, 0o644)
315 if err != nil {
316 t.Fatal(err)
317 }
318 if _, err := f.WriteString("{not valid json}\n"); err != nil {
319 _ = f.Close()
320 t.Fatal(err)
321 }
322 if err := f.Close(); err != nil {
323 t.Fatal(err)
324 }
325
326 d := newSessionTreeDialog()
327 if !d.OpenMessages([]provider.Message{treeTestMessage(provider.RoleUser, provider.TextBlock{Text: "keep"})}) {
328 t.Fatal("failed to seed dialog")
329 }
330 before := append([]sessionTreeItem(nil), d.items...)
331 if d.OpenSessionFamily(root, cwd, parent) {
332 t.Fatal("malformed family member passed preflight")
333 }
334 if !d.Active() || len(d.items) != len(before) || d.items[0].label != before[0].label {
335 t.Fatalf("preflight failure changed dialog: active=%v items=%+v", d.Active(), d.items)
336 }
337}
338
339func TestSanitizeSessionTreeTextPreservesTextAfterNonOSCExit(t *testing.T) {
340 input := "before\x1bM and \x1b(B after \x1b]0;window title\a tail"
341 if got, want := sanitizeSessionTreeText(input), "beforeM and (B after tail"; got != want {
342 t.Fatalf("sanitizeSessionTreeText = %q, want %q", got, want)
343 }
344}
345
346func TestSessionTreeRenderAndPagingAreWidthAndHeightSafe(t *testing.T) {
347 msgs := make([]provider.Message, 0, 20)
348 for i := 0; i < 20; i++ {
349 msgs = append(msgs, treeTestMessage(provider.RoleUser,
350 provider.TextBlock{Text: strings.Repeat("界", 4) + " long row"}))
351 }
352 d := newSessionTreeDialog()
353 d.MaxRows = 2
354 if !d.OpenMessages(msgs) {
355 t.Fatal("OpenMessages returned false")
356 }
357 if got := d.HandleKey(tui.Key{Kind: tui.KeyPageUp}); got.Select || d.cursor != 17 {
358 t.Fatalf("page up cursor = %d, action=%+v", d.cursor, got)
359 }
360 for _, width := range []int{0, 1, 7, 17} {
361 for _, line := range d.Render(tui.Theme{}, width) {
362 if got := sessionTreeANSIWidth(line); got > width {
363 t.Fatalf("width %d rendered %d cells: %q", width, got, line)
364 }
365 }
366 }
367 if d.viewTop < 0 || d.viewTop+d.MaxRows > len(d.items) {
368 t.Fatalf("invalid viewport top=%d rows=%d items=%d", d.viewTop, d.MaxRows, len(d.items))
369 }
370}
371
372func TestSessionTreeCheckoutRefreshesLastTurnUsageIncludingZero(t *testing.T) {
373 root := t.TempDir()
374 cwd := "/workspace/tree-usage"
375 path := newTreeTestSession(t, root, cwd, []provider.Message{
376 treeTestMessage(provider.RoleUser, provider.TextBlock{Text: "question"}),
377 treeTestMessage(provider.RoleAssistant, provider.TextBlock{Text: "answer"}),
378 })
379
380 initial := core.NewAgent(nil, "model", "", nil)
381 initial.SeedLastTurnUsage(provider.Usage{InputTokens: 99})
382 selected := core.NewAgent(nil, "model", "", nil)
383 selected.SetMessages(initial.Messages())
384 selected.SeedLastTurnUsage(provider.Usage{
385 InputTokens: 7,
386 CacheReadTokens: 11,
387 CacheWriteTokens: 13,
388 })
389
390 i := NewInteractive(InteractiveConfig{
391 Agent: initial,
392 CWD: cwd,
393 SessionsRoot: root,
394 CurrentSessionPath: func() string { return path },
395 })
396 i.cfg.LoadSession = func(string) error {
397 i.agent = selected
398 return nil
399 }
400
401 target := sessionTreeTarget{
402 SourcePath: path,
403 EffectiveIndex: 1,
404 SelectionBoundary: 2,
405 Role: provider.RoleAssistant,
406 }
407 i.applySessionTreeTarget(target, 1)
408 if i.lastCtxInput != 31 {
409 t.Fatalf("checked-out context usage = %d, want 31", i.lastCtxInput)
410 }
411
412 selected.SeedLastTurnUsage(provider.Usage{})
413 i.applySessionTreeTarget(target, 1)
414 if i.lastCtxInput != 0 {
415 t.Fatalf("checked-out zero usage left stale context value %d", i.lastCtxInput)
416 }
417}
418
419func treeTestMessage(role provider.Role, content ...provider.Content) provider.Message {
420 return provider.Message{Role: role, Content: content, Time: time.Unix(1, 0)}
421}
422
423func newTreeTestSession(t *testing.T, root, cwd string, msgs []provider.Message) string {
424 t.Helper()
425 sess, err := core.NewSession(root, cwd, "test", "model", "test")
426 if err != nil {
427 t.Fatal(err)
428 }
429 for _, msg := range msgs {
430 if err := sess.AppendMessage(msg); err != nil {
431 _ = sess.Close()
432 t.Fatal(err)
433 }
434 }
435 if err := sess.Close(); err != nil {
436 t.Fatal(err)
437 }
438 return sess.Path
439}