Agent-Native/runtime/session/tree_test.go

package session

import (
	"os"
	"path/filepath"
	"strings"
	"testing"

	"agent-native/runtime/llm"
)

func TestTreeForkCloneAndLeaf(t *testing.T) {
	dir := t.TempDir()
	s, err := Open(filepath.Join(dir, "a.jsonl"))
	if err != nil {
		t.Fatal(err)
	}
	if err := s.Append(llm.Message{Role: llm.RoleUser, Content: "one"}); err != nil {
		t.Fatal(err)
	}
	if err := s.Append(llm.Message{Role: llm.RoleAssistant, Content: "two"}); err != nil {
		t.Fatal(err)
	}
	entries, err := s.Entries()
	if err != nil || len(entries) != 2 || entries[1].ParentID != entries[0].ID {
		t.Fatalf("tree links: %+v %v", entries, err)
	}
	fork, err := s.Fork(filepath.Join(dir, "fork.jsonl"), entries[0].ID)
	if err != nil {
		t.Fatal(err)
	}
	if err := fork.Append(llm.Message{Role: llm.RoleUser, Content: "branch"}); err != nil {
		t.Fatal(err)
	}
	msgs, err := fork.Load()
	if err != nil || len(msgs) != 2 || msgs[1].Content != "branch" {
		t.Fatalf("fork path: %+v %v", msgs, err)
	}
	orig, _ := s.Load()
	if len(orig) != 2 || orig[1].Content != "two" {
		t.Fatalf("source branch changed: %+v", orig)
	}
	clone, err := s.Clone(filepath.Join(dir, "clone.jsonl"))
	if err != nil {
		t.Fatal(err)
	}
	cmsgs, _ := clone.Load()
	if len(cmsgs) != 2 || cmsgs[1].Content != "two" {
		t.Fatalf("clone: %+v", cmsgs)
	}
	if err := s.SetLeaf(entries[0].ID); err != nil {
		t.Fatal(err)
	}
	switched, _ := s.Load()
	if len(switched) != 1 || switched[0].Content != "one" {
		t.Fatalf("set leaf: %+v", switched)
	}
}

func TestLegacyTranscriptStillLoads(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "old.jsonl")
	body := "{\"role\":\"system\",\"content\":\"sys\"}\n{\"role\":\"user\",\"content\":\"hi\"}\n"
	if err := os.WriteFile(path, []byte(body), 0o644); err != nil {
		t.Fatal(err)
	}
	s, err := Open(path)
	if err != nil {
		t.Fatal(err)
	}
	msgs, err := s.Load()
	if err != nil || len(msgs) != 2 || msgs[1].Content != "hi" {
		t.Fatalf("legacy: %+v %v", msgs, err)
	}
	if err := s.Append(llm.Message{Role: llm.RoleAssistant, Content: "next"}); err != nil {
		t.Fatal(err)
	}
	msgs, err = s.Load()
	if err != nil || len(msgs) != 3 || msgs[2].Content != "next" {
		t.Fatalf("legacy append: %+v %v", msgs, err)
	}
}

func TestRepairTornTailBeforeAppend(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "torn.jsonl")
	// Valid record + torn final fragment (no closing brace / newline).
	body := "{\"role\":\"user\",\"content\":\"ok\"}\n{\"role\":\"assistant\",\"content\":\"tor"
	if err := os.WriteFile(path, []byte(body), 0o644); err != nil {
		t.Fatal(err)
	}
	s, err := Open(path)
	if err != nil {
		t.Fatal(err)
	}
	if err := s.Append(llm.Message{Role: llm.RoleAssistant, Content: "repaired"}); err != nil {
		t.Fatal(err)
	}
	msgs, err := s.Load()
	if err != nil {
		t.Fatalf("load after repair: %v", err)
	}
	if len(msgs) != 2 || msgs[0].Content != "ok" || msgs[1].Content != "repaired" {
		t.Fatalf("msgs after torn repair: %+v", msgs)
	}
	raw, _ := os.ReadFile(path)
	if !strings.HasSuffix(string(raw), "\n") {
		t.Fatal("file must end with newline after repair-append")
	}
	if strings.Contains(string(raw), "tor\n{") || strings.Contains(string(raw), "tor{") {
		t.Fatalf("torn fragment was concatenated: %s", raw)
	}
}

func TestRepairMissingNewlineBeforeAppend(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "nonewline.jsonl")
	// Complete JSON record without trailing newline.
	body := "{\"role\":\"user\",\"content\":\"hi\"}"
	if err := os.WriteFile(path, []byte(body), 0o644); err != nil {
		t.Fatal(err)
	}
	s, err := Open(path)
	if err != nil {
		t.Fatal(err)
	}
	if err := s.Append(llm.Message{Role: llm.RoleAssistant, Content: "next"}); err != nil {
		t.Fatal(err)
	}
	msgs, err := s.Load()
	if err != nil {
		t.Fatalf("load after newline repair: %v", err)
	}
	if len(msgs) != 2 || msgs[0].Content != "hi" || msgs[1].Content != "next" {
		t.Fatalf("msgs: %+v", msgs)
	}
	raw, _ := os.ReadFile(path)
	// Two separate JSON lines — not one concatenated object.
	lines := strings.Split(strings.TrimSuffix(string(raw), "\n"), "\n")
	if len(lines) != 2 {
		t.Fatalf("want 2 lines, got %d in %q", len(lines), raw)
	}
}

func TestMidFileCorruptionStillErrors(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "mid.jsonl")
	body := "{\"role\":\"user\",\"content\":\"a\"}\n{not-json}\n{\"role\":\"assistant\",\"content\":\"b\"}\n"
	if err := os.WriteFile(path, []byte(body), 0o644); err != nil {
		t.Fatal(err)
	}
	s, err := Open(path)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := s.Load(); err == nil {
		t.Fatal("mid-file corruption must be a hard error on load")
	}
	if err := s.Append(llm.Message{Role: llm.RoleUser, Content: "x"}); err == nil {
		t.Fatal("mid-file corruption must be a hard error on append")
	}
}

func TestRepairThenAppendThenLoad(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "roundtrip.jsonl")
	if err := os.WriteFile(path, []byte("{\"role\":\"user\",\"content\":\"one\"}\n{\"role\":\"as"), 0o644); err != nil {
		t.Fatal(err)
	}
	s, err := Open(path)
	if err != nil {
		t.Fatal(err)
	}
	if err := s.Append(llm.Message{Role: llm.RoleAssistant, Content: "two"}); err != nil {
		t.Fatal(err)
	}
	if err := s.Append(llm.Message{Role: llm.RoleUser, Content: "three"}); err != nil {
		t.Fatal(err)
	}
	msgs, err := s.Load()
	if err != nil || len(msgs) != 3 {
		t.Fatalf("roundtrip load: %+v %v", msgs, err)
	}
	if msgs[0].Content != "one" || msgs[1].Content != "two" || msgs[2].Content != "three" {
		t.Fatalf("contents: %+v", msgs)
	}
}

func TestCatalogContinueResumeAndHTML(t *testing.T) {
	dir := t.TempDir()
	s, err := Open(filepath.Join(dir, "sess.jsonl"))
	if err != nil {
		t.Fatal(err)
	}
	if err := s.SetCWD("/work"); err != nil {
		t.Fatal(err)
	}
	if err := s.SetName("notes"); err != nil {
		t.Fatal(err)
	}
	if err := s.Append(llm.Message{Role: llm.RoleUser, Content: "hello"}); err != nil {
		t.Fatal(err)
	}
	got, ok, err := Continue(dir, "/work")
	if err != nil || !ok || got.Path() != s.Path() {
		t.Fatalf("continue: %v %v %v", got, ok, err)
	}
	resumed, err := Resume(dir, "notes")
	if err != nil || resumed.Name() != "notes" {
		t.Fatalf("resume: %v %v", resumed, err)
	}
	htmlPath := filepath.Join(dir, "out.html")
	if err := s.ExportHTML(htmlPath); err != nil {
		t.Fatal(err)
	}
	raw, _ := os.ReadFile(htmlPath)
	if !strings.Contains(string(raw), "hello") || !strings.Contains(string(raw), "notes") {
		t.Fatalf("html: %s", raw)
	}
	if _, err := s.AppendEntry(Entry{Type: TypeLabel, Label: "checkpoint"}); err != nil {
		t.Fatal(err)
	}
	entries, _ := s.Entries()
	var labels int
	for _, e := range entries {
		if e.Type == TypeLabel && e.Label == "checkpoint" {
			labels++
		}
	}
	if labels != 1 {
		t.Fatalf("label missing: %+v", entries)
	}
}