fathom/internal/tracetree/tracetree_test.go

package tracetree

import (
	"os"
	"testing"

	"git.z.gd/zzg/fathom/internal/ingest"
	"git.z.gd/zzg/fathom/internal/model"
)

func span(id, parent, name string, start, end int64, attrs model.Attrs) model.Span {
	return model.Span{
		TraceID: "t1", SpanID: id, ParentSpanID: parent, Name: name,
		StartTS: start, EndTS: end, Attrs: attrs, SampleRate: 1,
	}
}

func TestBuildTree(t *testing.T) {
	spans := []model.Span{
		span("root", "", "interaction", 1000, 9000, nil),
		span("a", "root", "llm_request", 2000, 3000, nil),
		span("b", "root", "tool", 4000, 6000, nil),
		span("b1", "b", "tool.execution", 4500, 5500, nil),
	}
	tr := NewCorrelator().Build("t1", spans, nil)

	if len(tr.Roots) != 1 || tr.Roots[0].Span.SpanID != "root" {
		t.Fatalf("expected a single root named root, got %d roots", len(tr.Roots))
	}
	if tr.StartTS != 1000 || tr.EndTS != 9000 {
		t.Errorf("trace bounds = %d..%d, want 1000..9000", tr.StartTS, tr.EndTS)
	}

	flat := tr.Flatten()
	var order []string
	for _, n := range flat {
		order = append(order, n.Span.SpanID)
	}
	want := []string{"root", "a", "b", "b1"}
	for i := range want {
		if i >= len(order) || order[i] != want[i] {
			t.Fatalf("display order = %v, want %v", order, want)
		}
	}

	// Depth and offset drive the waterfall; both must be relative to the trace.
	byID := map[string]*Node{}
	for _, n := range flat {
		byID[n.Span.SpanID] = n
	}
	if byID["b1"].Depth != 2 {
		t.Errorf("b1 depth = %d, want 2", byID["b1"].Depth)
	}
	if byID["b"].OffsetMicr != 3000 {
		t.Errorf("b offset = %d, want 3000 (4000 - trace start 1000)", byID["b"].OffsetMicr)
	}
	if byID["b"].DurMicros != 2000 {
		t.Errorf("b duration = %d, want 2000", byID["b"].DurMicros)
	}
}

func TestOrphanBecomesRoot(t *testing.T) {
	// The exporter died before flushing the parent. The child must still show.
	spans := []model.Span{span("child", "missing-parent", "tool", 1000, 2000, nil)}
	tr := NewCorrelator().Build("t1", spans, nil)

	if len(tr.Roots) != 1 {
		t.Fatalf("orphan produced %d roots, want 1", len(tr.Roots))
	}
	if tr.OrphanCount != 1 {
		t.Errorf("OrphanCount = %d, want 1 — a truncated trace must say so", tr.OrphanCount)
	}
}

func TestCyclesDoNotHang(t *testing.T) {
	// A trace id collision can fabricate a cycle. Terminating with a slightly
	// wrong tree beats an unresponsive server.
	spans := []model.Span{
		span("a", "b", "x", 1000, 2000, nil),
		span("b", "a", "y", 1000, 2000, nil),
		span("self", "self", "z", 1000, 2000, nil),
	}
	done := make(chan *Trace, 1)
	go func() { done <- NewCorrelator().Build("t1", spans, nil) }()

	tr := <-done // a hang here fails the test by timeout, which is the point
	if got := len(tr.Flatten()); got != 3 {
		t.Errorf("flattened %d nodes, want all 3 spans exactly once", got)
	}
}

func TestLogsJoinByToolUseID(t *testing.T) {
	spans := []model.Span{
		span("root", "", "interaction", 1000, 9000, nil),
		span("tool", "root", "claude_code.tool", 4000, 6000,
			model.Attrs{"tool_use_id": model.StringValue("toolu_A")}),
		span("other", "root", "claude_code.tool", 7000, 8000,
			model.Attrs{"tool_use_id": model.StringValue("toolu_B")}),
	}
	// Both records carry the ROOT span id, as real emitters send them.
	logs := []model.LogRecord{
		{TraceID: "t1", SpanID: "root", EventName: "tool_result", TS: 5000,
			Attrs: model.Attrs{"tool_use_id": model.StringValue("toolu_A")}},
		{TraceID: "t1", SpanID: "root", EventName: "tool_result", TS: 7500,
			Attrs: model.Attrs{"tool_use_id": model.StringValue("toolu_B")}},
	}
	tr := NewCorrelator().Build("t1", spans, logs)

	byID := map[string]*Node{}
	for _, n := range tr.Flatten() {
		byID[n.Span.SpanID] = n
	}
	if len(byID["tool"].Logs) != 1 {
		t.Errorf("span tool got %d logs, want exactly its own", len(byID["tool"].Logs))
	}
	if len(byID["other"].Logs) != 1 {
		t.Errorf("span other got %d logs, want exactly its own", len(byID["other"].Logs))
	}
	// The failure mode this whole design exists to avoid.
	if len(byID["root"].Logs) != 0 {
		t.Errorf("root absorbed %d log records; the span_id join has crept back in", len(byID["root"].Logs))
	}
}

func TestLogsJoinByRequestID(t *testing.T) {
	spans := []model.Span{
		span("root", "", "interaction", 1000, 9000, nil),
		span("llm", "root", "claude_code.llm_request", 2000, 3000,
			model.Attrs{"request_id": model.StringValue("req_1")}),
	}
	logs := []model.LogRecord{
		{TraceID: "t1", SpanID: "root", EventName: "assistant_response", TS: 2500,
			Attrs: model.Attrs{"request_id": model.StringValue("req_1"),
				"response": model.StringValue("the answer")}},
	}
	tr := NewCorrelator().Build("t1", spans, logs)

	byID := map[string]*Node{}
	for _, n := range tr.Flatten() {
		byID[n.Span.SpanID] = n
	}
	if len(byID["llm"].Logs) != 1 {
		t.Fatalf("llm span got %d logs, want 1", len(byID["llm"].Logs))
	}
	if v := byID["llm"].Logs[0].Attrs["response"]; v.Str != "the answer" {
		t.Errorf("response text did not survive correlation: %q", v.Str)
	}
}

func TestParentPreferredOverChildOnSharedKey(t *testing.T) {
	// Claude Code emits `tool` and its child `tool.execution` with the same
	// tool_use_id. The result belongs on the one a reader calls "the tool call".
	spans := []model.Span{
		span("root", "", "claude_code.interaction", 1000, 9000, nil),
		span("tool", "root", "claude_code.tool", 4000, 6000,
			model.Attrs{"tool_use_id": model.StringValue("toolu_A")}),
		span("exec", "tool", "claude_code.tool.execution", 4100, 5900,
			model.Attrs{"tool_use_id": model.StringValue("toolu_A")}),
	}
	logs := []model.LogRecord{
		{TraceID: "t1", SpanID: "root", EventName: "tool_result", TS: 5000,
			Attrs: model.Attrs{"tool_use_id": model.StringValue("toolu_A")}},
	}
	tr := NewCorrelator().Build("t1", spans, logs)

	byID := map[string]*Node{}
	for _, n := range tr.Flatten() {
		byID[n.Span.SpanID] = n
	}
	if len(byID["tool"].Logs) != 1 {
		t.Errorf("tool_result landed on %q instead of the parent tool span",
			nodeWithLogs(byID))
	}
}

func nodeWithLogs(byID map[string]*Node) string {
	for id, n := range byID {
		if len(n.Logs) > 0 {
			return id
		}
	}
	return "(nothing)"
}

func TestKeylessLogPrefersItsOwnSpanID(t *testing.T) {
	// A user prompt has no identity key but does carry the root's span id, and
	// it belongs to the run as a whole. Timestamp containment would file it
	// under whichever model call started at the same moment, which is wrong.
	spans := []model.Span{
		span("root", "", "interaction", 1000, 9000, nil),
		span("inner", "root", "llm_request", 1000, 3000, nil),
	}
	logs := []model.LogRecord{
		{TraceID: "t1", SpanID: "root", EventName: "user_prompt", TS: 1000,
			Attrs: model.Attrs{"prompt": model.StringValue("hello")}},
	}
	tr := NewCorrelator().Build("t1", spans, logs)

	byID := map[string]*Node{}
	for _, n := range tr.Flatten() {
		byID[n.Span.SpanID] = n
	}
	if len(byID["root"].Logs) != 1 {
		t.Errorf("user_prompt should sit on the interaction it names, not on %q", nodeWithLogs(byID))
	}
	if len(byID["inner"].Logs) != 0 {
		t.Errorf("user_prompt leaked onto the model call that shares its start time")
	}
}

func TestKeylessLogWithUnknownSpanFallsBackToContainment(t *testing.T) {
	// When the span id names nothing in this trace, timestamp containment is
	// the only signal left, and is better than dropping the record.
	spans := []model.Span{
		span("root", "", "interaction", 1000, 9000, nil),
		span("inner", "root", "llm_request", 2000, 3000, nil),
	}
	logs := []model.LogRecord{
		{TraceID: "t1", SpanID: "never-ingested", EventName: "note", TS: 2500,
			Attrs: model.Attrs{"x": model.StringValue("y")}},
	}
	tr := NewCorrelator().Build("t1", spans, logs)

	byID := map[string]*Node{}
	for _, n := range tr.Flatten() {
		byID[n.Span.SpanID] = n
	}
	if len(byID["inner"].Logs) != 1 {
		t.Errorf("record with an unknown span id did not fall back to the containing span")
	}
}

func TestTraceStatusErrorIfAnySpanErrored(t *testing.T) {
	spans := []model.Span{
		span("root", "", "interaction", 1000, 9000, nil),
		span("bad", "root", "tool", 2000, 3000, nil),
	}
	spans[1].Status = model.StatusError
	if got := NewCorrelator().Build("t1", spans, nil).Status; got != "error" {
		t.Errorf("trace status = %q, want error when a child span failed", got)
	}
}

// TestAgainstCapturedClaudeCode is the end-to-end check on the real payloads:
// decode what Claude Code actually sent, correlate it, and assert the result is
// a readable transcript rather than everything piled on the root.
func TestAgainstCapturedClaudeCode(t *testing.T) {
	d := ingest.NewDecoder()

	traceBytes, err := os.ReadFile("../ingest/testdata/claude_code_traces.json")
	if err != nil {
		t.Fatal(err)
	}
	logBytes, err := os.ReadFile("../ingest/testdata/claude_code_logs_content.json")
	if err != nil {
		t.Fatal(err)
	}
	spans, err := d.DecodeTraces(traceBytes)
	if err != nil {
		t.Fatal(err)
	}
	logs, err := d.DecodeLogs(logBytes)
	if err != nil {
		t.Fatal(err)
	}
	if len(spans) == 0 || len(logs) == 0 {
		t.Fatalf("fixtures decoded to %d spans and %d logs", len(spans), len(logs))
	}

	tr := NewCorrelator().Build(spans[0].TraceID, spans, logs)

	if len(tr.Roots) == 0 {
		t.Fatal("no roots in a trace that has a clear single root")
	}
	flat := tr.Flatten()
	if len(flat) != len(spans) {
		t.Errorf("flattened %d nodes from %d spans; some were lost or duplicated", len(flat), len(spans))
	}

	// The real assertion: content must be distributed across spans, not heaped
	// on the root. With the naive span_id join every record lands on one node.
	withLogs := 0
	var rootLogs int
	for _, n := range flat {
		if len(n.Logs) > 0 {
			withLogs++
		}
		if n.Span.ParentSpanID == "" {
			rootLogs += len(n.Logs)
		}
	}
	if withLogs < 2 {
		t.Errorf("only %d span(s) received log records; correlation collapsed onto one node", withLogs)
	}

	// And specifically: a tool span must have got its own tool_result.
	var toolMatched bool
	for _, n := range flat {
		if n.Span.Kind != model.KindToolCall {
			continue
		}
		for _, l := range n.Logs {
			if l.EventName == "tool_result" {
				toolMatched = true
			}
		}
	}
	if !toolMatched {
		t.Errorf("no tool span received a tool_result record (%d spans had logs, %d on root)", withLogs, rootLogs)
	}
}