ollie/cmd/olliesrv/internal/agent/chat_test.go

package agent

import (
	"context"
	"strings"
	"sync"
	"testing"
	"time"

	"ollie/format"
)

// testAgent creates a minimal Agent for chat/block testing.
func testAgent(sessionID, agentID string) *Agent {
	ag := &Agent{
		sessionID: sessionID,
		id:        agentID,
		state:     "idle",
	}
	ag.signalCh = make(chan struct{})
	ag.chatSignalCh = make(chan struct{})
	ag.chatCond = sync.NewCond(ag.chatMu.RLocker())
	return ag
}

func TestNextBlockID_Deterministic(t *testing.T) {
	ag1 := testAgent("session1", "agent1")
	ag2 := testAgent("session1", "agent1")

	// Same session/agent should produce same sequence
	id1a := ag1.NextBlockID()
	id1b := ag1.NextBlockID()
	id2a := ag2.NextBlockID()
	id2b := ag2.NextBlockID()

	if id1a != id2a {
		t.Errorf("first block IDs don't match: %s vs %s", id1a, id2a)
	}
	if id1b != id2b {
		t.Errorf("second block IDs don't match: %s vs %s", id1b, id2b)
	}
	if id1a == id1b {
		t.Errorf("sequential block IDs should differ: %s == %s", id1a, id1b)
	}
}

func TestNextBlockID_DifferentAgents(t *testing.T) {
	ag1 := testAgent("session1", "agent1")
	ag2 := testAgent("session1", "agent2")

	id1 := ag1.NextBlockID()
	id2 := ag2.NextBlockID()

	if id1 == id2 {
		t.Errorf("different agents should have different block IDs: %s == %s", id1, id2)
	}
}

func TestNextBlockID_DifferentSessions(t *testing.T) {
	ag1 := testAgent("session1", "agent1")
	ag2 := testAgent("session2", "agent1")

	id1 := ag1.NextBlockID()
	id2 := ag2.NextBlockID()

	if id1 == id2 {
		t.Errorf("different sessions should have different block IDs: %s == %s", id1, id2)
	}
}

func TestNextBlockID_Length(t *testing.T) {
	ag := testAgent("test-session", "test-agent")
	id := ag.NextBlockID()

	// Block IDs are 8 hex chars (from sha256[:8])
	if len(id) != 8 {
		t.Errorf("block ID length should be 8, got %d: %s", len(id), id)
	}
	// Should be valid hex
	for _, c := range id {
		if !((c >= '0' && c <= '9') || (c >= 'a' && c <= 'f')) {
			t.Errorf("block ID contains non-hex char: %s", id)
			break
		}
	}
}

func TestAppendBlock_JSONL(t *testing.T) {
	ag := testAgent("session1", "agent1")

	ag.AppendBlock(format.Block{Role: "user", ID: "abc12345", Content: "Hello"})
	ag.AppendBlock(format.Block{Role: "assistant", ID: "def67890", Content: "Hi there"})
	ag.AppendBlock(format.Block{Role: "tool", ID: "11223344", Content: "output"})

	ag.ChatMu().RLock()
	raw := ag.RawLog()
	ag.ChatMu().RUnlock()

	// Should contain valid JSONL
	if len(raw) == 0 {
		t.Fatal("raw log should not be empty")
	}

	// Parse and verify
	lines := 0
	for _, line := range splitLines(raw) {
		if len(line) == 0 {
			continue
		}
		lines++
		b, err := format.UnmarshalBlock(line)
		if err != nil {
			t.Errorf("failed to parse JSONL line: %v", err)
		}
		if b.ID == "" {
			t.Error("block should have ID")
		}
	}
	if lines != 3 {
		t.Errorf("expected 3 blocks, got %d", lines)
	}
}

func TestBlockByID(t *testing.T) {
	ag := testAgent("session1", "agent1")

	ag.AppendBlock(format.Block{Role: "user", ID: "abc12345", Content: "Hello"})
	ag.AppendBlock(format.Block{Role: "assistant", ID: "def67890", Content: "Hi there"})

	// Find existing block
	block, found := ag.BlockByID("def67890")
	if !found {
		t.Fatal("block should be found")
	}
	if block.Role != "assistant" || block.Content != "Hi there" {
		t.Errorf("wrong block content: %+v", block)
	}

	// Not found
	_, found = ag.BlockByID("notfound")
	if found {
		t.Error("nonexistent block should not be found")
	}

	// Empty ID
	_, found = ag.BlockByID("")
	if found {
		t.Error("empty block ID should not be found")
	}
}

func TestBlockByID_SkipsPartial(t *testing.T) {
	ag := testAgent("session1", "agent1")

	// Add partial block first
	ag.AppendBlock(format.Block{Role: "assistant", ID: "abc12345", Content: "H", Partial: true})
	// Then final
	ag.AppendBlock(format.Block{Role: "assistant", ID: "abc12345", Content: "Hello"})

	block, found := ag.BlockByID("abc12345")
	if !found {
		t.Fatal("block should be found")
	}
	if block.Content != "Hello" {
		t.Errorf("should find final block, not partial: %+v", block)
	}
}

func TestAppendBlock_SignalsBroadcast(t *testing.T) {
	ag := testAgent("session1", "agent1")

	signalCh := ag.ChatSignal()

	// Append block
	ag.AppendBlock(format.Block{Role: "user", ID: "test", Content: "hi"})

	// Signal channel should be closed
	select {
	case <-signalCh:
		// Good - channel closed
	default:
		t.Error("chat signal should have fired")
	}

	// New signal channel should be open
	newSignalCh := ag.ChatSignal()
	select {
	case <-newSignalCh:
		t.Error("new signal channel should not be closed yet")
	default:
		// Good - channel open
	}
}

func TestRawLogRead_Streaming(t *testing.T) {
	ag := testAgent("session1", "agent1")

	// Initial read with no base - should return empty and current offset
	data, base, err := ag.RawLogRead("")
	if err != nil {
		t.Fatalf("RawLogRead error: %v", err)
	}
	if len(data) != 0 {
		t.Errorf("initial read should be empty, got: %s", data)
	}

	// Append a block
	ag.AppendBlock(format.Block{Role: "user", ID: "test", Content: "hello"})

	// Read with previous base - should get new data
	data, newBase, err := ag.RawLogRead(base)
	if err != nil {
		t.Fatalf("RawLogRead error: %v", err)
	}
	if len(data) == 0 {
		t.Error("should have new data")
	}

	// Verify it's valid JSON
	b, err := format.UnmarshalBlock(data[:len(data)-1]) // strip trailing newline
	if err != nil {
		t.Errorf("data should be valid JSONL: %v", err)
	}
	if b.Content != "hello" {
		t.Errorf("wrong content: %s", b.Content)
	}

	// Read again with new base - should be empty
	data, _, err = ag.RawLogRead(newBase)
	if err != nil {
		t.Fatalf("RawLogRead error: %v", err)
	}
	if len(data) != 0 {
		t.Errorf("no new data, should be empty, got: %s", data)
	}
}

func TestTextLog_RendersHumanReadable(t *testing.T) {
	ag := testAgent("session1", "agent1")

	ag.AppendBlock(format.Block{Role: "user", ID: "a", Content: "Hello"})
	ag.AppendBlock(format.Block{Role: "assistant", ID: "b", Content: "Hi there"})
	ag.AppendBlock(format.Block{Role: "call", ID: "c", Name: "shell", Content: `{"cmd":"ls"}`})

	ag.ChatMu().RLock()
	text := string(ag.TextLog())
	ag.ChatMu().RUnlock()

	// Should have user and assistant content
	if !contains(text, "Hello") {
		t.Error("should contain user message")
	}
	if !contains(text, "Hi there") {
		t.Error("should contain assistant message")
	}
	// Tool calls are now surfaced in the rendered log (name + args).
	if !contains(text, "shell") {
		t.Error("call should show tool name")
	}
	if !contains(text, "cmd") {
		t.Error("call should show tool args")
	}
}

// TestTextLog_IncrementalPartialNoDuplication verifies that assistant partials
// stream into textLog incrementally and that finalization does not duplicate the
// already-streamed body (only the trailing separator is appended).
func TestTextLog_IncrementalPartialNoDuplication(t *testing.T) {
	ag := testAgent("session1", "agent1")
	id := ag.NextBlockID()

	// Stream partials: "Hel" -> "Hello" -> "Hello world"
	ag.SetPartial(format.Block{Role: "assistant", ID: id, Content: "Hel", Partial: true})
	if got := string(ag.TextLog()); got != "Hel" {
		t.Fatalf("after first partial, textLog = %q, want %q", got, "Hel")
	}
	ag.SetPartial(format.Block{Role: "assistant", ID: id, Content: "Hello", Partial: true})
	if got := string(ag.TextLog()); got != "Hello" {
		t.Fatalf("after second partial, textLog = %q, want %q", got, "Hello")
	}
	ag.SetPartial(format.Block{Role: "assistant", ID: id, Content: "Hello world", Partial: true})

	// Finalize.
	ag.AppendBlock(format.Block{Role: "assistant", ID: id, Content: "Hello world"})

	got := string(ag.TextLog())
	if got != "Hello world\n\n" {
		t.Errorf("textLog = %q, want %q (body must not be duplicated)", got, "Hello world\n\n")
	}
	if strings.Count(got, "Hello world") != 1 {
		t.Errorf("body appears %d times, want 1: %q", strings.Count(got, "Hello world"), got)
	}
}

// TestRawLogStream_NoLostWakeup verifies that a finalized block appended while
// a stream reader is between its data check and its wait is still delivered,
// without requiring a second event. This is the regression test for the TUI /
// last-message-lost race.
func TestRawLogStream_NoLostWakeup(t *testing.T) {
	ag := testAgent("session1", "agent1")
	ag.AppendBlock(format.Block{Role: "user", ID: "a", Content: "hi"})

	// Catch up: fresh open returns the existing block and a sync base.
	data, base, err := ag.RawLogStream(context.Background(), "")
	if err != nil {
		t.Fatalf("RawLogStream error: %v", err)
	}
	if len(data) == 0 {
		t.Fatal("expected initial history")
	}

	// Reader is now caught up. Append concurrently while it blocks; the
	// stream must return the new block promptly (not hang).
	go func() {
		ag.AppendBlock(format.Block{Role: "assistant", ID: "b", Content: "last message"})
	}()

	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	data, _, err = ag.RawLogStream(ctx, base)
	if err != nil {
		t.Fatalf("RawLogStream error: %v", err)
	}
	if ctx.Err() != nil {
		t.Fatal("stream blocked: final block was lost (lost wakeup)")
	}
	if !contains(string(data), "last message") {
		t.Errorf("expected final block delivered, got %q", data)
	}
}

// Helper to split JSONL
func splitLines(data []byte) [][]byte {
	var lines [][]byte
	start := 0
	for i, b := range data {
		if b == '\n' {
			lines = append(lines, data[start:i])
			start = i + 1
		}
	}
	if start < len(data) {
		lines = append(lines, data[start:])
	}
	return lines
}

func contains(s, substr string) bool {
	return len(s) >= len(substr) && (s == substr || len(substr) == 0 ||
		(len(s) > 0 && len(substr) > 0 && findSubstring(s, substr)))
}

func findSubstring(s, substr string) bool {
	for i := 0; i <= len(s)-len(substr); i++ {
		if s[i:i+len(substr)] == substr {
			return true
		}
	}
	return false
}