feat(trace): 深追踪（Span 生命周期/父子链/JSONL recorder，支持 llm/tool/state 等类别）
This commit is contained in:
Aska committed
1 parent
011d25a53d
commit
e33b005612
4 files changed
+661
No files matched your search
@@ -0,0 +1,333 @@
|
||||
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)
|
||||
}
|
||||
if sp.Duration <= 0 {
|
||||
t.Fatalf("Duration = %v, want > 0", sp.Duration)
|
||||
}
|
||||
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)
|
||||
}
|
||||
if childRec.Duration <= 0 {
|
||||
t.Fatalf("childRec.Duration = %v, want > 0", childRec.Duration)
|
||||
}
|
||||
if childRec.Props["model"] != "gpt-test" {
|
||||
t.Fatalf("childRec.Props = %+v, want model key", childRec.Props)
|
||||
}
|
||||
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])
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user