Cyrene_v2/internal/runtime/trace/trace_test.go

package trace

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"strings"
	"sync"
	"testing"
	"time"
)

// decodeLines reads every JSON object emitted as JSONL into recs, failing the
// test if any line is malformed.
func decodeLines(t *testing.T, data []byte) []SpanRecord {
	t.Helper()
	sc := bufio.NewScanner(bytes.NewReader(data))
	var recs []SpanRecord
	for sc.Scan() {
		line := sc.Bytes()
		if len(bytes.TrimSpace(line)) == 0 {
			continue
		}
		var r SpanRecord
		if err := json.Unmarshal(line, &r); err != nil {
			t.Fatalf("decode span: %v", err)
		}
		recs = append(recs, r)
	}
	if err := sc.Err(); err != nil {
		t.Fatalf("scan spans: %v", err)
	}
	return recs
}

func TestSpan_Lifecycle(t *testing.T) {
	t.Parallel()
	var buf bytes.Buffer
	tr := New(NewJSONRecorder(&buf))

	ctx, sp := tr.StartSpan(context.Background(), "llm-call", SpanKindLLM)
	if sp.ID == "" {
		t.Fatal("span ID is empty")
	}
	if sp.Kind != SpanKindLLM || sp.Name != "llm-call" {
		t.Fatalf("unexpected (name, kind) = (%q, %q)", sp.Name, sp.Kind)
	}
	if sp.StartedAt.IsZero() {
		t.Fatalf("StartedAt is zero, want set at creation")
	}
	if sp.Status != "" || sp.Duration != 0 || !sp.EndedAt.IsZero() {
		t.Fatalf("Status=%q Duration=%v EndedAt=%v, want zero before Finish", sp.Status, sp.Duration, sp.EndedAt)
	}
	if got := SpanFromContext(ctx); got != sp {
		t.Fatalf("SpanFromContext = %p, want %p", got, sp)
	}

	before := time.Now()
	if err := sp.SetProp("model", "gpt-test"); err != nil {
		t.Fatalf("SetProp: %v", err)
	}
	if err := sp.Record(Event{Type: "token", Props: map[string]any{"n": 1}}); err != nil {
		t.Fatalf("Record: %v", err)
	}

	sp.Finish(SpanStatusOK)

	if sp.Status != SpanStatusOK {
		t.Fatalf("Status = %q, want ok", sp.Status)
	}
	if sp.EndedAt.Before(before) {
		t.Fatalf("EndedAt %v is before snapshot %v", sp.EndedAt, before)
	}
	// The monotonic clock may be too coarse to tick between StartedAt and
	// EndedAt, so a sub-tick span legitimately reads Duration == 0. Assert the
	// platform-independent invariant instead: Duration is non-negative and is
	// exactly the delta of the two stamps.
	if sp.Duration < 0 {
		t.Fatalf("Duration = %v, want >= 0", sp.Duration)
	}
	if want := sp.EndedAt.Sub(sp.StartedAt); sp.Duration != want {
		t.Fatalf("Duration = %v, want EndedAt-StartedAt = %v", sp.Duration, want)
	}
	if len(sp.Events) != 1 || sp.Events[0].Type != "token" {
		t.Fatalf("Events = %+v, want one 'token' event", sp.Events)
	}
	if sp.Props["model"] != "gpt-test" {
		t.Fatalf("Props = %+v, want model key", sp.Props)
	}

	recs := decodeLines(t, buf.Bytes())
	if len(recs) != 1 {
		t.Fatalf("emitted %d records, want 1", len(recs))
	}
}

func TestSpan_RecordStampsTime(t *testing.T) {
	t.Parallel()
	tr := New(nil)
	_, sp := tr.StartSpan(context.Background(), "x", SpanKindTool)

	if err := sp.Record(Event{Type: "retry"}); err != nil {
		t.Fatalf("Record: %v", err)
	}
	if sp.Events[0].Time.IsZero() {
		t.Fatalf("event time was not stamped: %v", sp.Events[0].Time)
	}
}

func TestSpan_FinishTwiceIgnored(t *testing.T) {
	t.Parallel()
	var buf bytes.Buffer
	tr := New(NewJSONRecorder(&buf))
	_, sp := tr.StartSpan(context.Background(), "x", SpanKindGuardrail)

	sp.Finish(SpanStatusError)
	first := sp.Duration
	sp.Finish(SpanStatusOK)

	if sp.Status != SpanStatusError {
		t.Fatalf("Status = %q, want error (second Finish ignored)", sp.Status)
	}
	if sp.Duration != first {
		t.Fatalf("Duration changed from %v to %v on second Finish", first, sp.Duration)
	}
	if recs := decodeLines(t, buf.Bytes()); len(recs) != 1 {
		t.Fatalf("emitted %d records, want 1", len(recs))
	}
}

func TestSpan_MutationAfterFinish(t *testing.T) {
	t.Parallel()
	tr := New(nil)
	_, sp := tr.StartSpan(context.Background(), "x", SpanKindMemory)
	sp.Finish(SpanStatusOK)

	if err := sp.Record(Event{Type: "late"}); err != ErrSpanFinished {
		t.Fatalf("Record after Finish = %v, want ErrSpanFinished", err)
	}
	if err := sp.SetProp("late", true); err != ErrSpanFinished {
		t.Fatalf("SetProp after Finish = %v, want ErrSpanFinished", err)
	}
	if len(sp.Events) != 0 || len(sp.Props) != 0 {
		t.Fatalf("closed span mutated: events=%d props=%d", len(sp.Events), len(sp.Props))
	}
}

func TestTracer_ParentChildChain(t *testing.T) {
	t.Parallel()
	tr := New(nil)

	ctx, root := tr.StartSpan(context.Background(), "run", SpanKindOrchestration)
	ctx, child := tr.StartSpan(ctx, "llm", SpanKindLLM)

	if child.ParentID != root.ID {
		t.Fatalf("child.ParentID = %q, want root %q", child.ParentID, root.ID)
	}
	if root.ParentID != "" {
		t.Fatalf("root.ParentID = %q, want empty", root.ParentID)
	}
	if got := SpanFromContext(ctx); got != child {
		t.Fatalf("SpanFromContext = %p, want child %p", got, child)
	}

	// A sibling span is chained to the same parent via the parent's context.
	_, sibling := tr.StartSpan(context.Background(), "tool", SpanKindTool)
	if sibling.ParentID != "" {
		t.Fatalf("sibling.ParentID = %q, want empty (no parent in fresh ctx)", sibling.ParentID)
	}
}

func TestTracer_TwoLevels(t *testing.T) {
	t.Parallel()
	tr := New(nil)
	ctx, root := tr.StartSpan(context.Background(), "run", SpanKindOrchestration)
	ctx, mid := tr.StartSpan(ctx, "plan", SpanKindOrchestration)
	ctx, leaf := tr.StartSpan(ctx, "llm", SpanKindLLM)

	if leaf.ParentID != mid.ID {
		t.Fatalf("leaf.ParentID = %q, want mid %q", leaf.ParentID, mid.ID)
	}
	if mid.ParentID != root.ID {
		t.Fatalf("mid.ParentID = %q, want root %q", mid.ParentID, root.ID)
	}
}

func TestTracer_NilAndEmptyContext(t *testing.T) {
	t.Parallel()
	tr := New(nil)

	if got := SpanFromContext(nil); got != nil {
		t.Fatalf("SpanFromContext(nil) = %p, want nil", got)
	}
	if got := SpanFromContext(context.Background()); got != nil {
		t.Fatalf("SpanFromContext(bg) = %p, want nil", got)
	}

	if _, sp := tr.StartSpan(nil, "x", SpanKindStateTransition); sp == nil {
		t.Fatal("StartSpan(nil ctx) returned nil span")
	}
}

func TestTracer_UniqueIDs(t *testing.T) {
	t.Parallel()
	tr := New(nil)
	ids := make(map[string]bool)
	for i := 0; i < 100; i++ {
		_, sp := tr.StartSpan(context.Background(), "x", SpanKindTool)
		if sp.ID == "" {
			t.Fatal("empty span ID")
		}
		if ids[sp.ID] {
			t.Fatalf("duplicate span ID %q", sp.ID)
		}
		ids[sp.ID] = true
	}
}

func TestJSONRecorder_Output(t *testing.T) {
	t.Parallel()
	var buf bytes.Buffer
	tr := New(NewJSONRecorder(&buf))

	ctx, root := tr.StartSpan(context.Background(), "run", SpanKindOrchestration)
	_, child := tr.StartSpan(ctx, "llm", SpanKindLLM)
	_ = child.SetProp("model", "gpt-test")
	_ = child.Record(Event{Type: "prompt", Props: map[string]any{"tokens": 12}})
	child.Finish(SpanStatusOK)
	root.Finish(SpanStatusOK)

	recs := decodeLines(t, buf.Bytes())
	if len(recs) != 2 {
		t.Fatalf("emitted %d records, want 2: %+v", len(recs), recs)
	}

	var childRec, rootRec SpanRecord
	for _, r := range recs {
		if r.Kind == SpanKindLLM {
			childRec = r
		} else {
			rootRec = r
		}
	}
	if childRec.ParentID != rootRec.ID {
		t.Fatalf("childRec.ParentID = %q, want root %q", childRec.ParentID, rootRec.ID)
	}
	if childRec.Status != SpanStatusOK {
		t.Fatalf("childRec.Status = %q, want ok", childRec.Status)
	}
	// As in TestSpan_Lifecycle: a sub-tick span reads Duration == 0 on platforms
	// with a coarse monotonic clock, so assert the recorded delta invariant.
	if childRec.Duration < 0 {
		t.Fatalf("childRec.Duration = %v, want >= 0", childRec.Duration)
	}
	if want := childRec.EndedAt.Sub(childRec.StartedAt); childRec.Duration != want {
		t.Fatalf("childRec.Duration = %v, want EndedAt-StartedAt = %v", childRec.Duration, want)
	}
	if _, ok := childRec.Props["model"]; ok {
		t.Fatalf("childRec.Props = %+v，表外键 model 被原样落盘了（写侧只投影白名单）", childRec.Props)
	}
	omitted, ok := childRec.Props[omittedPropKey].(map[string]any)
	if !ok {
		t.Fatalf("childRec.Props = %+v，表外键没有折进 %q", childRec.Props, omittedPropKey)
	}
	if _, ok := omitted["model"]; !ok {
		t.Fatalf("omitted = %+v，want 一条 model 的 Digest（脱敏后仍要留诊断线索）", omitted)
	}
	if len(childRec.Events) != 1 || childRec.Events[0].Type != "prompt" {
		t.Fatalf("childRec.Events = %+v, want one 'prompt' event", childRec.Events)
	}

	// Ensure lines are well-formed JSON objects (not interleaved), by re-decoding.
	for _, r := range recs {
		if r.ID == "" {
			t.Fatal("record with empty ID was emitted")
		}
	}
}

func TestTracer_Concurrent(t *testing.T) {
	t.Parallel()
	var buf bytes.Buffer
	tr := New(NewJSONRecorder(&buf))

	const workers = 64
	var wg sync.WaitGroup
	for i := 0; i < workers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			ctx, parent := tr.StartSpan(context.Background(), "run", SpanKindOrchestration)
			for j := 0; j < 5; j++ {
				_, child := tr.StartSpan(ctx, "llm", SpanKindLLM)
				if child.ParentID != parent.ID {
					t.Errorf("child.ParentID = %q, want parent %q", child.ParentID, parent.ID)
					return
				}
				_ = child.SetProp("worker", "w")
				_ = child.Record(Event{Type: "done"})
				child.Finish(SpanStatusOK)
			}
			parent.Finish(SpanStatusOK)
		}()
	}
	wg.Wait()

	recs := decodeLines(t, buf.Bytes())
	// Each worker emits one parent + 5 children.
	if len(recs) != workers*6 {
		t.Fatalf("emitted %d records, want %d", len(recs), workers*6)
	}

	// No two records may share an ID.
	seen := make(map[string]bool, len(recs))
	for _, r := range recs {
		if seen[r.ID] {
			t.Fatalf("duplicate ID %q emitted concurrently", r.ID)
		}
		seen[r.ID] = true
	}
}

func TestTracer_NilWriterSpansCreateOnly(t *testing.T) {
	t.Parallel()
	tr := New(nil)
	_, sp := tr.StartSpan(context.Background(), "x", SpanKindSubagent)
	sp.Finish(SpanStatusOK)
	if sp.Status != SpanStatusOK {
		t.Fatalf("Status = %q, want ok", sp.Status)
	}
}

func TestSpanKind_Values(t *testing.T) {
	t.Parallel()
	want := []SpanKind{
		SpanKindLLM, SpanKindTool, SpanKindOrchestration, SpanKindStateTransition,
		SpanKindSubagent, SpanKindMemory, SpanKindGuardrail,
	}
	got := []string{
		"llm_call", "tool_call", "orchestration", "state_transition",
		"subagent", "memory", "guardrail",
	}
	for i := range want {
		if string(want[i]) != got[i] {
			t.Fatalf("SpanKind[%d] = %q, want %q", i, want[i], got[i])
		}
	}
	if !strings.Contains(string(want[0]), "llm_call") {
		t.Fatalf("kind string mismatch: %s", want[0])
	}
}