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