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