fathom/internal/tracetree/tracetree.go

// Package tracetree reconstructs a single trace: its span tree, and the log
// records that belong to each span.
//
// This is deliberately not built on the query engine. Opening a trace is the
// day-to-day reason Fathom exists, and it wants a different shape of work than
// an aggregate scan: find one trace id, rebuild a tree, fetch blobs for the one
// span someone clicked. Expressing that through a SQL-shaped planner would make
// both worse.
package tracetree

import (
	"sort"
	"strings"

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

// Node is a span plus its children and correlated log records.
type Node struct {
	Span     model.Span        `json:"span"`
	Logs     []model.LogRecord `json:"logs,omitempty"`
	Children []*Node           `json:"children,omitempty"`

	// Depth and the timing fields are precomputed because the UI needs them for
	// every row and recomputing them per render is the kind of thing that makes
	// a 2000-span trace feel slow.
	Depth      int   `json:"depth"`
	OffsetMicr int64 `json:"offset_micros"` // from the trace's start
	DurMicros  int64 `json:"duration_micros"`
}

// Trace is a fully reconstructed trace.
type Trace struct {
	TraceID   string  `json:"trace_id"`
	Roots     []*Node `json:"roots"`
	StartTS   int64   `json:"start_ts"`
	EndTS     int64   `json:"end_ts"`
	SpanCount int     `json:"span_count"`
	Status    string  `json:"status"`

	// Orphans are spans whose parent was never ingested. They are surfaced
	// rather than dropped: a missing parent usually means the exporter was
	// still batching when the process exited, and silently hiding those spans
	// would make a trace look complete when it is not.
	OrphanCount int `json:"orphan_count"`
}

// Correlator attaches log records to the spans they describe.
//
// The naive approach — join on span_id — does not work, and it fails in a way
// that looks like success. Real emitters stamp content records with the id of
// the root span for the whole run, not the span the content belongs to. Joining
// on span_id therefore piles every prompt, tool result, and response onto the
// root node and leaves every interesting span empty. Verified against captured
// Claude Code output, where all 17 correlated records carried the same root
// span id.
//
// What does work is a typed join on identity keys that both sides already
// carry: a tool call id links a tool-result record to the tool span that ran
// it, and a request id links a model response to the LLM span that produced it.
// Those keys are configurable because their names differ between
// instrumentations.
type Correlator struct {
	// JoinKeys are attribute names tried in order. The first key present on
	// both a log record and a span wins.
	JoinKeys []string
}

// DefaultJoinKeys covers Claude Code and the OpenTelemetry GenAI conventions.
// Order matters: the most specific identifier is tried first, so a record
// carrying both a tool call id and a request id attaches to the tool.
func DefaultJoinKeys() []string {
	return []string{
		"tool_use_id",         // Claude Code tool spans
		"gen_ai.tool.call.id", // OTel GenAI conventions
		"request_id",          // Claude Code LLM spans
		"client_request_id",   // Claude Code, client-side correlation
		"gen_ai.response.id",  // OTel GenAI conventions
	}
}

func NewCorrelator() *Correlator {
	return &Correlator{JoinKeys: DefaultJoinKeys()}
}

// Build assembles a trace from its spans and log records.
func (c *Correlator) Build(traceID string, spans []model.Span, logs []model.LogRecord) *Trace {
	t := &Trace{TraceID: traceID, SpanCount: len(spans)}
	if len(spans) == 0 {
		return t
	}

	// Sorting by start time makes sibling order deterministic and puts the
	// waterfall in the order things actually happened.
	sort.SliceStable(spans, func(i, j int) bool {
		if spans[i].StartTS != spans[j].StartTS {
			return spans[i].StartTS < spans[j].StartTS
		}
		return spans[i].SpanID < spans[j].SpanID
	})

	nodes := make(map[string]*Node, len(spans))
	for i := range spans {
		s := spans[i]
		nodes[s.SpanID] = &Node{Span: s, DurMicros: s.DurationMicros()}
	}

	c.attachLogs(nodes, spans, logs)

	// Link children to parents. A span whose parent is absent becomes a root so
	// that it still renders; an incomplete trace is more useful than a blank one.
	for i := range spans {
		s := &spans[i]
		n := nodes[s.SpanID]
		parent, ok := nodes[s.ParentSpanID]
		switch {
		case s.ParentSpanID == "":
			t.Roots = append(t.Roots, n)
		case !ok:
			t.OrphanCount++
			t.Roots = append(t.Roots, n)
		case parent == n:
			// A span claiming itself as parent would loop forever below.
			t.OrphanCount++
			t.Roots = append(t.Roots, n)
		default:
			parent.Children = append(parent.Children, n)
		}
	}

	// A cycle among spans would make the walk below non-terminating. Emitters
	// should never produce one, but a trace id collision between two runs can
	// fabricate one, and an unresponsive server is a worse failure than a
	// slightly wrong tree.
	breakCycles(t, nodes)

	t.StartTS, t.EndTS = spanBounds(spans)
	t.Status = traceStatus(spans)

	for _, root := range t.Roots {
		assignLayout(root, 0, t.StartTS)
	}
	sort.SliceStable(t.Roots, func(i, j int) bool {
		return t.Roots[i].Span.StartTS < t.Roots[j].Span.StartTS
	})
	return t
}

// attachLogs performs the typed join described on Correlator.
func (c *Correlator) attachLogs(nodes map[string]*Node, spans []model.Span, logs []model.LogRecord) {
	// Index spans by every join key they carry.
	index := make(map[string]map[string]*Node, len(c.JoinKeys))
	for _, key := range c.JoinKeys {
		index[key] = make(map[string]*Node)
	}
	for i := range spans {
		n := nodes[spans[i].SpanID]
		for _, key := range c.JoinKeys {
			if v, ok := spans[i].Attrs[key]; ok && v.Str != "" {
				// Several spans can share a key: Claude Code emits both a
				// `tool` span and its child `tool.execution` with the same
				// tool_use_id. Prefer the shallower one — the parent is the
				// span a reader thinks of as "the tool call".
				if prev, taken := index[key][v.Str]; !taken || isAncestorName(spans[i].Name, prev.Span.Name) {
					index[key][v.Str] = n
				}
			}
		}
	}

	var unattached []model.LogRecord
	for _, rec := range logs {
		if n := c.match(index, rec); n != nil {
			n.Logs = append(n.Logs, rec)
			continue
		}
		unattached = append(unattached, rec)
	}

	// Records with no identity key — a user prompt, a session-level event — fall
	// back to their span_id.
	//
	// This is not the span_id join this type warns about, and the difference
	// matters. A record carrying a typed key is *about* a specific tool call or
	// model request, and its span_id is the root's, so the key must win. A
	// record carrying no key is about the run as a whole, and for those the
	// root span_id it was stamped with is exactly right. Preferring timestamp
	// containment here instead would file the user's prompt under whichever
	// model call happened to start at the same microsecond, which is how this
	// was caught: the prompt rendered under llm_request rather than under the
	// interaction that owns it.
	for _, rec := range unattached {
		if n, ok := nodes[rec.SpanID]; ok {
			n.Logs = append(n.Logs, rec)
			continue
		}
		if n := containingSpan(nodes, rec); n != nil {
			n.Logs = append(n.Logs, rec)
		}
	}

	for _, n := range nodes {
		sort.SliceStable(n.Logs, func(i, j int) bool { return n.Logs[i].TS < n.Logs[j].TS })
	}
}

func (c *Correlator) match(index map[string]map[string]*Node, rec model.LogRecord) *Node {
	for _, key := range c.JoinKeys {
		v, ok := rec.Attrs[key]
		if !ok || v.Str == "" {
			continue
		}
		if n, ok := index[key][v.Str]; ok {
			return n
		}
	}
	return nil
}

// isAncestorName reports whether a is the shallower of two related span names,
// using the dotted-prefix convention emitters use for nested spans
// ("claude_code.tool" vs "claude_code.tool.execution").
func isAncestorName(a, b string) bool {
	return len(a) < len(b) && strings.HasPrefix(b, a+".")
}

// containingSpan finds the deepest span whose time window contains the record,
// so an uncorrelated event lands as close to its real origin as the timestamps
// allow.
func containingSpan(nodes map[string]*Node, rec model.LogRecord) *Node {
	var best *Node
	var bestDur int64
	for _, n := range nodes {
		if rec.TS < n.Span.StartTS || rec.TS > n.Span.EndTS {
			continue
		}
		d := n.Span.DurationMicros()
		if best == nil || d < bestDur {
			best, bestDur = n, d
		}
	}
	return best
}

// breakCycles ensures the tree is walkable. Any node not reachable from a root
// is promoted to a root, which turns a cycle into a set of subtrees.
func breakCycles(t *Trace, nodes map[string]*Node) {
	seen := make(map[*Node]bool, len(nodes))
	var walk func(*Node) bool
	walk = func(n *Node) bool {
		if seen[n] {
			return false // already placed, or a cycle
		}
		seen[n] = true
		kept := n.Children[:0]
		for _, ch := range n.Children {
			if walk(ch) {
				kept = append(kept, ch)
			}
		}
		n.Children = kept
		return true
	}
	for _, r := range t.Roots {
		walk(r)
	}
	for _, n := range nodes {
		if !seen[n] {
			seen[n] = true
			n.Children = nil
			t.Roots = append(t.Roots, n)
			t.OrphanCount++
		}
	}
}

func spanBounds(spans []model.Span) (start, end int64) {
	for i := range spans {
		s := &spans[i]
		if start == 0 || s.StartTS < start {
			start = s.StartTS
		}
		if s.EndTS > end {
			end = s.EndTS
		}
	}
	return start, end
}

// traceStatus is error if any span errored. A trace where one tool call failed
// is a trace worth noticing, even when the run recovered.
func traceStatus(spans []model.Span) string {
	for i := range spans {
		if spans[i].Status == model.StatusError {
			return "error"
		}
	}
	return "ok"
}

// assignLayout precomputes depth and waterfall offsets, and orders children by
// start time.
func assignLayout(n *Node, depth int, traceStart int64) {
	n.Depth = depth
	n.OffsetMicr = n.Span.StartTS - traceStart
	if n.OffsetMicr < 0 {
		n.OffsetMicr = 0
	}
	sort.SliceStable(n.Children, func(i, j int) bool {
		return n.Children[i].Span.StartTS < n.Children[j].Span.StartTS
	})
	for _, ch := range n.Children {
		assignLayout(ch, depth+1, traceStart)
	}
}

// Flatten returns the tree in depth-first display order, which is the order the
// UI renders rows in.
func (t *Trace) Flatten() []*Node {
	var out []*Node
	var walk func(*Node)
	walk = func(n *Node) {
		out = append(out, n)
		for _, ch := range n.Children {
			walk(ch)
		}
	}
	for _, r := range t.Roots {
		walk(r)
	}
	return out
}