ollie/fs/replay_test.go

package fs

import (
	"encoding/json"
	"strings"
	"testing"

	"ollie/backend"
	"ollie/format"
)

// checkLogStructure verifies that the replayed log has balanced blocks
// (every [[[header]]] has a matching [[[end]]]) with no nesting.
func checkLogStructure(t *testing.T, log []byte) {
	t.Helper()
	depth := 0
	for i, line := range strings.Split(string(log), "\n") {
		trimmed := strings.TrimSpace(line)
		if trimmed == "" {
			continue
		}
		if format.IsEndMarker(trimmed) {
			depth--
			if depth < 0 {
				t.Fatalf("line %d: unexpected [[[end]]] with no open block", i+1)
			}
		} else if bh := format.ParseBlockHeader(trimmed); bh != nil {
			if depth > 0 {
				t.Fatalf("line %d: nested block [[[%s:%s]]] at depth %d", i+1, bh.Role, bh.Name, depth)
			}
			depth++
		}
	}
	if depth != 0 {
		t.Fatalf("log ended with %d unclosed block(s)", depth)
	}
}

func TestReplayMessagesToLog_assistantWithToolCalls(t *testing.T) {
	al := NewAgentLog("test")

	messages := []backend.Message{
		{Role: "user", Content: "hello"},
		{
			Role:    "assistant",
			ID:      "resp-123",
			Content: "Let me check that.",
			ToolCalls: []backend.ToolCall{
				{ID: "tc1", Name: "tool_load", Arguments: json.RawMessage(`{"name":"shell"}`)},
			},
		},
		{Role: "tool", Content: "ok", ToolCallID: "tc1"},
		{
			Role:    "assistant",
			ID:      "resp-456",
			Content: "Done!",
		},
	}

	replayMessagesToLog(al, messages)
	checkLogStructure(t, al.Log())
}

func TestReplayMessagesToLog_multipleToolCalls(t *testing.T) {
	al := NewAgentLog("test")

	messages := []backend.Message{
		{Role: "user", Content: "do stuff"},
		{
			Role:    "assistant",
			ID:      "resp-abc",
			Content: "I'll call two tools.",
			ToolCalls: []backend.ToolCall{
				{ID: "tc1", Name: "file_read", Arguments: json.RawMessage(`{"path":"/tmp/x"}`)},
				{ID: "tc2", Name: "file_grep", Arguments: json.RawMessage(`{"pattern":"foo"}`)},
			},
		},
		{Role: "tool", Content: "file contents", ToolCallID: "tc1"},
		{Role: "tool", Content: "grep results", ToolCallID: "tc2"},
		{
			Role:    "assistant",
			ID:      "resp-def",
			Content: "All done.",
		},
	}

	replayMessagesToLog(al, messages)
	checkLogStructure(t, al.Log())
}

func TestReplayMessagesToLog_assistantNoContent(t *testing.T) {
	al := NewAgentLog("test")

	messages := []backend.Message{
		{Role: "user", Content: "run something"},
		{
			Role: "assistant",
			ID:   "resp-000",
			ToolCalls: []backend.ToolCall{
				{ID: "tc1", Name: "shell", Arguments: json.RawMessage(`{"cmd":"echo hi"}`)},
			},
		},
		{Role: "tool", Content: "hi\n", ToolCallID: "tc1"},
	}

	replayMessagesToLog(al, messages)
	checkLogStructure(t, al.Log())
}

func TestReplayMessagesToLog_consecutiveAssistantAfterTool(t *testing.T) {
	al := NewAgentLog("test")

	messages := []backend.Message{
		{Role: "user", Content: "fix the bug"},
		{
			Role:    "assistant",
			ID:      "resp-aaa",
			Content: "Let me look at the code.",
			ToolCalls: []backend.ToolCall{
				{ID: "call_1", Name: "file_read", Arguments: json.RawMessage(`{"path":"/tmp/main.go"}`)},
			},
		},
		{Role: "tool", Content: "package main\n", ToolCallID: "call_1"},
		{
			Role:    "assistant",
			ID:      "resp-bbb",
			Content: "I see the issue. Fixed it.",
		},
	}

	replayMessagesToLog(al, messages)
	checkLogStructure(t, al.Log())
}

func TestReplayMessagesToLog_manyTurns(t *testing.T) {
	al := NewAgentLog("test")

	messages := []backend.Message{
		{Role: "user", Content: "hello"},
		{Role: "assistant", ID: "r1", Content: "Hi! How can I help?"},
		{Role: "user", Content: "read a file"},
		{
			Role:    "assistant",
			ID:      "r2",
			Content: "Sure.",
			ToolCalls: []backend.ToolCall{
				{ID: "c1", Name: "file_read", Arguments: json.RawMessage(`{"path":"/tmp/x"}`)},
			},
		},
		{Role: "tool", Content: "contents", ToolCallID: "c1"},
		{
			Role:    "assistant",
			ID:      "r3",
			Content: "Here's the file content.",
		},
		{Role: "user", Content: "now grep for foo"},
		{
			Role: "assistant",
			ID:   "r4",
			ToolCalls: []backend.ToolCall{
				{ID: "c2", Name: "file_grep", Arguments: json.RawMessage(`{"pattern":"foo"}`)},
			},
		},
		{Role: "tool", Content: "line 5: foo bar", ToolCallID: "c2"},
		{Role: "assistant", ID: "r5", Content: "Found foo on line 5."},
	}

	replayMessagesToLog(al, messages)
	checkLogStructure(t, al.Log())
}