Agent-Native/runtime/session/r2_regression_test.go

package session

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

	"agent-native/runtime/llm"
)

func TestR2CompactKeepsToolPair(t *testing.T) {
	s, err := Open(filepath.Join(t.TempDir(), "s.jsonl"))
	if err != nil {
		t.Fatal(err)
	}
	msgs := []llm.Message{
		{Role: llm.RoleUser, Content: "ask"},
		{Role: llm.RoleAssistant, Content: "", ToolCalls: []llm.ToolCall{{ID: "c", Name: "read"}}},
		{Role: llm.RoleTool, ToolCallID: "c", Content: "file-body"},
	}
	for _, m := range msgs {
		if err := s.Append(m); err != nil {
			t.Fatal(err)
		}
	}
	ok, err := s.Compact(1, func(messages []llm.Message) (string, error) {
		return "summary", nil
	})
	if err != nil || !ok {
		t.Fatalf("compact: %v %v", ok, err)
	}
	loaded, err := s.Load()
	if err != nil {
		t.Fatal(err)
	}
	var sawTool bool
	for _, m := range loaded {
		if m.Role == llm.RoleTool && m.Content == "file-body" {
			sawTool = true
		}
	}
	if !sawTool {
		t.Fatalf("tool result dropped: %+v", loaded)
	}
}

func TestR2CompactRejectsAppendDuringSummary(t *testing.T) {
	s, err := Open(filepath.Join(t.TempDir(), "s.jsonl"))
	if err != nil {
		t.Fatal(err)
	}
	for _, text := range []string{"a", "b", "c"} {
		if err := s.Append(llm.Message{Role: llm.RoleUser, Content: text}); err != nil {
			t.Fatal(err)
		}
	}
	_, err = s.Compact(1, func(messages []llm.Message) (string, error) {
		return "summary", s.Append(llm.Message{Role: llm.RoleUser, Content: "raced"})
	})
	if err == nil || !strings.Contains(err.Error(), "compaction conflict") {
		t.Fatalf("want conflict, got %v", err)
	}
}

func TestCompactRetriesThenCommits(t *testing.T) {
	s, err := Open(filepath.Join(t.TempDir(), "s.jsonl"))
	if err != nil {
		t.Fatal(err)
	}
	for _, text := range []string{"a", "b", "c"} {
		if err := s.Append(llm.Message{Role: llm.RoleUser, Content: text}); err != nil {
			t.Fatal(err)
		}
	}
	raced := false
	ok, err := s.Compact(1, func(messages []llm.Message) (string, error) {
		if !raced {
			raced = true
			if err := s.Append(llm.Message{Role: llm.RoleUser, Content: "raced"}); err != nil {
				return "", err
			}
			return "", fmt.Errorf("compaction conflict: session changed during summary")
		}
		return "summary", nil
	})
	if err != nil || !ok {
		t.Fatalf("retry compact: %v %v", ok, err)
	}
	loaded, err := s.Load()
	if err != nil {
		t.Fatal(err)
	}
	var saw bool
	for _, m := range loaded {
		if m.Content == "raced" || m.Content == "summary" {
			saw = true
		}
	}
	if !saw {
		t.Fatalf("retried summary missing: %+v", loaded)
	}
}

func TestCompactLeavesNoOrphanToolCall(t *testing.T) {
	s, err := Open(filepath.Join(t.TempDir(), "s.jsonl"))
	if err != nil {
		t.Fatal(err)
	}
	msgs := []llm.Message{
		{Role: llm.RoleUser, Content: "先看"},
		{Role: llm.RoleAssistant, ToolCalls: []llm.ToolCall{{ID: "old", Name: "read_file"}}},
		{Role: llm.RoleTool, ToolCallID: "old", Name: "read_file", Content: "旧"},
		{Role: llm.RoleUser, Content: "再改"},
		{Role: llm.RoleAssistant, ToolCalls: []llm.ToolCall{{ID: "a", Name: "read_file"}, {ID: "b", Name: "grep"}}},
		{Role: llm.RoleTool, ToolCallID: "a", Name: "read_file", Content: "甲"},
		{Role: llm.RoleTool, ToolCallID: "b", Name: "grep", Content: "乙"},
	}
	for _, m := range msgs {
		if err := s.Append(m); err != nil {
			t.Fatal(err)
		}
	}
	ok, err := s.Compact(2, func(messages []llm.Message) (string, error) {
		return "摘要", nil
	})
	if err != nil || !ok {
		t.Fatalf("compact: %v %v", ok, err)
	}
	loaded, err := s.Load()
	if err != nil {
		t.Fatal(err)
	}
	assertToolPairs(t, loaded)
}

func assertToolPairs(t *testing.T, msgs []llm.Message) {
	t.Helper()
	calls := map[string]int{}
	for _, m := range msgs {
		if m.Role != llm.RoleAssistant {
			continue
		}
		for _, c := range m.ToolCalls {
			calls[c.ID]++
		}
	}
	for _, m := range msgs {
		if m.Role != llm.RoleTool {
			continue
		}
		if calls[m.ToolCallID] == 0 {
			t.Fatalf("orphan tool result %s in %+v", m.ToolCallID, msgs)
		}
		calls[m.ToolCallID]--
	}
	for id, n := range calls {
		if n != 0 {
			t.Fatalf("orphan tool call %s in %+v", id, msgs)
		}
	}
}

func TestR2AppendVisibleAfterReopen(t *testing.T) {
	path := filepath.Join(t.TempDir(), "s.jsonl")
	s, err := Open(path)
	if err != nil {
		t.Fatal(err)
	}
	if err := s.Append(llm.Message{Role: llm.RoleUser, Content: "durable"}); err != nil {
		t.Fatal(err)
	}
	again, err := Open(path)
	if err != nil {
		t.Fatal(err)
	}
	msgs, err := again.Load()
	if err != nil {
		t.Fatal(err)
	}
	if len(msgs) != 1 || msgs[0].Content != "durable" {
		t.Fatalf("reopen: %+v", msgs)
	}
}

// ---- N2 WP1：alignToolTail 边界直接单测（原仅经 Store.Compact 间接覆盖）----

func tailMsgs() []llm.Message {
	call := func(id string) llm.ToolCall { return llm.ToolCall{ID: id, Name: "read", Arguments: "{}"} }
	return []llm.Message{
		{Role: llm.RoleUser, Content: "u1"},                                                  // 0
		{Role: llm.RoleAssistant, ToolCalls: []llm.ToolCall{call("a"), call("b")}},          // 1 多调用 assistant
		{Role: llm.RoleTool, ToolCallID: "a", Content: "ra"},                                 // 2
		{Role: llm.RoleTool, ToolCallID: "b", Content: "rb"},                                 // 3
		{Role: llm.RoleUser, Content: "u2"},                                                  // 4
		{Role: llm.RoleAssistant, ToolCalls: []llm.ToolCall{call("c")}},                     // 5
		{Role: llm.RoleTool, ToolCallID: "c", Content: "rc"},                                 // 6
	}
}

func TestAlignToolTailEdges(t *testing.T) {
	msgs := tailMsgs()
	// tail 恰落在多调用 assistant 的结果组内 → 回退到该 assistant。
	if got := alignToolTail(msgs, 3); got != 1 {
		t.Fatalf("落在多调用结果组内：got %d want 1", got)
	}
	// tail 落在结果组边界外（user）→ 不动。
	if got := alignToolTail(msgs, 4); got != 4 {
		t.Fatalf("落在 user 上：got %d want 4", got)
	}
	// 结尾整组是 tool 角色 → 回退到其 assistant。
	if got := alignToolTail(msgs, 6); got != 5 {
		t.Fatalf("落在末尾结果上：got %d want 5", got)
	}
	// 从头起全是 tool 角色（退化形态）→ 回退到 0，不越界。
	allTool := []llm.Message{
		{Role: llm.RoleTool, ToolCallID: "x", Content: "1"},
		{Role: llm.RoleTool, ToolCallID: "y", Content: "2"},
	}
	if got := alignToolTail(allTool, 1); got != 0 {
		t.Fatalf("全 tool：got %d want 0", got)
	}
}