diff --git a/client.go b/client.go index 089d9e7..f8d526f 100644 --- a/client.go +++ b/client.go @@ -266,3 +266,31 @@ func (c *Client) Permalink(span oteltrace.Span) string { } return link } + +// Export serializes the span into a string that can be passed to another process +// and used with ContextWithExportedSpan to create child spans there. The format +// is compatible with the Braintrust JS/Python span.export() output. +// +// Returns an empty string and logs a warning on error. +func (c *Client) Export(span oteltrace.Span) string { + s, err := bttrace.Export(span) + if err != nil { + c.logger.Warn("could not export span", "error", err) + return "" + } + return s +} + +// ContextWithExportedSpan returns a context that carries the exported span as the +// remote parent. Spans started with this context will be children of that span. +// The exported string should be from Export(span) or span.export() (JS/Python). +// +// Returns the original context and logs a warning on error. +func (c *Client) ContextWithExportedSpan(ctx context.Context, exported string) context.Context { + out, err := bttrace.ContextWithExportedSpan(ctx, exported) + if err != nil { + c.logger.Warn("could not set context from exported span", "error", err) + return ctx + } + return out +} diff --git a/examples/distributed-tracing/main.go b/examples/distributed-tracing/main.go index 50a21da..5963741 100644 --- a/examples/distributed-tracing/main.go +++ b/examples/distributed-tracing/main.go @@ -1,9 +1,14 @@ -// Package main demonstrates distributed tracing using W3C baggage propagation. +// Package main demonstrates distributed tracing using W3C baggage propagation +// and Braintrust span export (similar to span.export() in the JS/Python SDK). // // This example shows how trace context propagates across service boundaries // via W3C baggage. A parent span encodes context to headers, and a child span // extracts it (simulated without an actual HTTP server). // +// Alternatively, you can use trace.Export(span) to get a serialized string and +// trace.ContextWithExportedSpan(ctx, exported) on the remote side to attach +// children to that span (e.g. when passing the parent in a message or custom header). +// // To run this example: // // export BRAINTRUST_API_KEY="your-api-key" @@ -20,6 +25,7 @@ import ( sdktrace "go.opentelemetry.io/otel/sdk/trace" "github.com/braintrustdata/braintrust-sdk-go" + bttrace "github.com/braintrustdata/braintrust-sdk-go/trace" ) func main() { @@ -59,6 +65,11 @@ func main() { // Call remote service (simulates crossing service boundary) simulateHTTPRequest(headers) + // Alternative: pass exported span string (like JS/Python span.export()) + if exported, err := bttrace.Export(parentSpan); err == nil { + simulateHTTPRequestWithExportedSpan(exported) + } + // Flush all spans if err := tp.ForceFlush(context.Background()); err != nil { log.Printf("Failed to flush spans: %v", err) @@ -78,3 +89,19 @@ func simulateHTTPRequest(headers map[string]string) { _, span := tracer.Start(ctx, "remote-service.handle-request") defer span.End() } + +// simulateHTTPRequestWithExportedSpan simulates a remote service that receives +// an exported span string (e.g. from a message or custom header) and creates +// a child span. This is similar to using span.export() in the JS/Python SDK. +func simulateHTTPRequestWithExportedSpan(exported string) { + tracer := otel.Tracer("examples/distributed-tracing") + + ctx, err := bttrace.ContextWithExportedSpan(context.Background(), exported) + if err != nil { + log.Printf("ContextWithExportedSpan: %v", err) + return + } + + _, span := tracer.Start(ctx, "remote-service.handle-request-export") + defer span.End() +} diff --git a/examples/export-span/main.go b/examples/export-span/main.go new file mode 100644 index 0000000..adb6ae6 --- /dev/null +++ b/examples/export-span/main.go @@ -0,0 +1,68 @@ +// Package main demonstrates linking distributed traces using span export. +// +// Service A starts a root span and exports it (like span.export() in JS/Python). +// The exported string is passed to Service B (e.g. via message queue or HTTP header). +// Service B uses ContextWithExportedSpan to create a child span in the same trace. +// +// To run: +// +// export BRAINTRUST_API_KEY="your-api-key" +// go run examples/export-span/main.go +package main + +import ( + "context" + "fmt" + "log" + + sdktrace "go.opentelemetry.io/otel/sdk/trace" + oteltrace "go.opentelemetry.io/otel/trace" + + "github.com/braintrustdata/braintrust-sdk-go" + bttrace "github.com/braintrustdata/braintrust-sdk-go/trace" +) + +func main() { + tp := sdktrace.NewTracerProvider() + defer tp.Shutdown(context.Background()) //nolint:errcheck + + bt, err := braintrust.New(tp, + braintrust.WithProject("go-sdk-examples"), + braintrust.WithBlockingLogin(true), + ) + if err != nil { + log.Fatalf("Failed to initialize Braintrust: %v", err) + } + + tracer := bt.Tracer("export-span-example") + + // --- Service A: create root span and export it --- + _, rootSpan := tracer.Start(context.Background(), "service-a.request") + exported := bt.Export(rootSpan) + if exported == "" { + log.Fatal("Export failed") + } + defer rootSpan.End() + + // Simulate sending the exported string to Service B (e.g. in a message or header) + // In a real setup: publish to a queue, put in HTTP header "X-Braintrust-Span", etc. + runServiceB(tracer, exported) + + if err := tp.ForceFlush(context.Background()); err != nil { + log.Printf("Flush: %v", err) + } + + fmt.Printf("\nView trace: %s\n", bt.Permalink(rootSpan)) +} + +// runServiceB simulates a separate service that receives the exported span +// and creates a child span in the same trace. +func runServiceB(tracer oteltrace.Tracer, exported string) { + ctx, err := bttrace.ContextWithExportedSpan(context.Background(), exported) + if err != nil { + log.Fatalf("ContextWithExportedSpan: %v", err) + } + + _, childSpan := tracer.Start(ctx, "service-b.process") + childSpan.End() +} diff --git a/trace/span_components.go b/trace/span_components.go new file mode 100644 index 0000000..8734d21 --- /dev/null +++ b/trace/span_components.go @@ -0,0 +1,279 @@ +// Package trace provides span component encoding/decoding for distributed tracing. +// The format is compatible with the Braintrust JS/Python SDK SpanComponents V4. +package trace + +import ( + "encoding/base64" + "encoding/hex" + "encoding/json" + "fmt" + "regexp" +) + +const encodingVersionV4 = 4 + +// Field IDs in the V4 binary format (must match JS/Python). +const ( + fieldObjectID = 1 + fieldRowID = 2 + fieldSpanID = 3 // 8-byte hex + fieldRootSpanID = 4 // 16-byte hex +) + +// SpanObjectTypeV3 matches the Braintrust backend object types. +type SpanObjectTypeV3 byte + +const ( + SpanObjectTypeExperiment SpanObjectTypeV3 = 1 + SpanObjectTypeProjectLogs SpanObjectTypeV3 = 2 + SpanObjectTypePlaygroundLogs SpanObjectTypeV3 = 3 +) + +// SpanComponents holds the decoded span identifiers from an exported span. +// It is compatible with the Braintrust SpanComponents V4 format. +type SpanComponents struct { + ObjectType SpanObjectTypeV3 + ObjectID string + + // Span identifiers (all set when exporting a specific span). + RowID string + SpanID string // 8-byte hex (16 chars) + RootSpanID string // 16-byte hex (32 chars), same as OTel trace ID +} + +// tryHexSpanID checks if s is 16 hex chars (8 bytes) and returns bytes or nil. +func tryHexSpanID(s string) []byte { + if len(s) != 16 { + return nil + } + b, err := hex.DecodeString(s) + if err != nil { + return nil + } + return b +} + +// tryHexTraceID checks if s is 32 hex chars (16 bytes) and returns bytes or nil. +func tryHexTraceID(s string) []byte { + if len(s) != 32 { + return nil + } + b, err := hex.DecodeString(s) + if err != nil { + return nil + } + return b +} + +// EncodeV4 serializes the span components to a base64 string (SpanComponents V4 format). +// The result can be passed to other processes and used with ContextWithExportedSpan. +func EncodeV4(c SpanComponents) (string, error) { + jsonObj := make(map[string]interface{}) + + // Header: version + object_type + buf := []byte{encodingVersionV4, byte(c.ObjectType)} + + var hexEntries [][]byte + + addHexField := func(val string, fieldID byte) { + if fieldID == fieldSpanID { + if b := tryHexSpanID(val); b != nil { + hexEntries = append(hexEntries, append([]byte{fieldID}, b...)) + return + } + } else if fieldID == fieldRootSpanID { + if b := tryHexTraceID(val); b != nil { + hexEntries = append(hexEntries, append([]byte{fieldID}, b...)) + return + } + } + // Non-hex or other field: put in JSON + name := map[byte]string{fieldObjectID: "object_id", fieldRowID: "row_id", fieldSpanID: "span_id", fieldRootSpanID: "root_span_id"}[fieldID] + jsonObj[name] = val + } + + if c.ObjectID != "" { + addHexField(c.ObjectID, fieldObjectID) + } + if c.RowID != "" { + addHexField(c.RowID, fieldRowID) + } + if c.SpanID != "" { + addHexField(c.SpanID, fieldSpanID) + } + if c.RootSpanID != "" { + addHexField(c.RootSpanID, fieldRootSpanID) + } + + if len(hexEntries) > 255 { + return "", fmt.Errorf("too many hex entries") + } + + buf = append(buf, byte(len(hexEntries))) + for _, e := range hexEntries { + buf = append(buf, e...) + } + + if len(jsonObj) > 0 { + jsonBytes, err := json.Marshal(jsonObj) + if err != nil { + return "", err + } + buf = append(buf, jsonBytes...) + } + + return base64.StdEncoding.EncodeToString(buf), nil +} + +// DecodeV4 parses a base64-encoded SpanComponents V4 string. +// Supports both V4 and older V3 payloads (decodes via V3 logic for version < 4). +func DecodeV4(s string) (SpanComponents, error) { + raw, err := base64.StdEncoding.DecodeString(s) + if err != nil { + return SpanComponents{}, fmt.Errorf("decode base64: %w", err) + } + if len(raw) < 3 { + return SpanComponents{}, fmt.Errorf("span components too short") + } + + version := raw[0] + if version < encodingVersionV4 { + return decodeV3Fallback(raw, s) + } + + // V4 format + out := SpanComponents{ + ObjectType: SpanObjectTypeV3(raw[1]), + } + numHex := raw[2] + offset := 3 + + for i := 0; i < int(numHex); i++ { + if offset >= len(raw) { + break + } + fieldID := raw[offset] + offset++ + switch fieldID { + case fieldSpanID: + if offset+8 > len(raw) { + return SpanComponents{}, fmt.Errorf("truncated span_id") + } + out.SpanID = hex.EncodeToString(raw[offset : offset+8]) + offset += 8 + case fieldRootSpanID: + if offset+16 > len(raw) { + return SpanComponents{}, fmt.Errorf("truncated root_span_id") + } + out.RootSpanID = hex.EncodeToString(raw[offset : offset+16]) + offset += 16 + case fieldObjectID, fieldRowID: + if offset+16 > len(raw) { + return SpanComponents{}, fmt.Errorf("truncated field %d", fieldID) + } + val := hex.EncodeToString(raw[offset : offset+16]) + offset += 16 + if fieldID == fieldObjectID { + out.ObjectID = val + } else { + out.RowID = val + } + default: + return SpanComponents{}, fmt.Errorf("unknown field id %d", fieldID) + } + } + + if offset < len(raw) { + var jsonObj map[string]interface{} + if err := json.Unmarshal(raw[offset:], &jsonObj); err != nil { + return SpanComponents{}, fmt.Errorf("decode json: %w", err) + } + for k, v := range jsonObj { + if s, ok := v.(string); ok { + switch k { + case "object_id": + out.ObjectID = s + case "row_id": + out.RowID = s + case "span_id": + out.SpanID = s + case "root_span_id": + out.RootSpanID = s + } + } + } + } + + return out, nil +} + +// uuidLike matches UUID format (with hyphens) or 32 hex chars. +var uuidLike = regexp.MustCompile(`^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$|^[0-9a-fA-F]{32}$`) + +// ParentFromComponents returns a Parent that can be used with SetParent so spans +// started under ContextWithExportedSpan are sent to the same Braintrust project/experiment. +// For PROJECT_LOGS, object_id is treated as project name unless it looks like a UUID. +func ParentFromComponents(c SpanComponents) (Parent, error) { + switch c.ObjectType { + case SpanObjectTypeExperiment: + return Parent{Type: ParentTypeExperimentID, ID: c.ObjectID}, nil + case SpanObjectTypeProjectLogs, SpanObjectTypePlaygroundLogs: + if uuidLike.MatchString(c.ObjectID) { + return Parent{Type: ParentTypeProjectID, ID: c.ObjectID}, nil + } + return Parent{Type: ParentTypeProjectName, ID: c.ObjectID}, nil + default: + return Parent{}, fmt.Errorf("unsupported object type %d", c.ObjectType) + } +} + +// decodeV3Fallback decodes V3-style encoding (version < 4) for interoperability. +// V3 uses 16-byte UUIDs for all fields; we only support reading object_type and object_id from the UUID section, +// and span_id/root_span_id if present as UUIDs. For simplicity we parse the minimal required fields. +func decodeV3Fallback(raw []byte, _ string) (SpanComponents, error) { + if len(raw) < 3 { + return SpanComponents{}, fmt.Errorf("v3 span components too short") + } + out := SpanComponents{ + ObjectType: SpanObjectTypeV3(raw[1]), + } + numUUID := raw[2] + offset := 3 + for i := 0; i < int(numUUID); i++ { + if offset+17 > len(raw) { + break + } + fieldID := raw[offset] + uuidBytes := raw[offset+1 : offset+17] + offset += 17 + // UUID bytes to hex string (32 chars) + val := hex.EncodeToString(uuidBytes) + switch fieldID { + case fieldObjectID: + out.ObjectID = val + case fieldRowID: + out.RowID = val + case fieldSpanID: + out.SpanID = val + case fieldRootSpanID: + out.RootSpanID = val + } + } + if offset < len(raw) { + var jsonObj map[string]interface{} + _ = json.Unmarshal(raw[offset:], &jsonObj) + if v, ok := jsonObj["object_id"].(string); ok && out.ObjectID == "" { + out.ObjectID = v + } + if v, ok := jsonObj["row_id"].(string); ok { + out.RowID = v + } + if v, ok := jsonObj["span_id"].(string); ok { + out.SpanID = v + } + if v, ok := jsonObj["root_span_id"].(string); ok { + out.RootSpanID = v + } + } + return out, nil +} diff --git a/trace/span_components_test.go b/trace/span_components_test.go new file mode 100644 index 0000000..e91716e --- /dev/null +++ b/trace/span_components_test.go @@ -0,0 +1,81 @@ +package trace + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestEncodeDecodeV4_RoundTrip(t *testing.T) { + c := SpanComponents{ + ObjectType: SpanObjectTypeProjectLogs, + ObjectID: "my-project", + RowID: "1234567890abcdef", + SpanID: "1234567890abcdef", + RootSpanID: "a1b2c3d4e5f6789012345678abcdef01", + } + + encoded, err := EncodeV4(c) + require.NoError(t, err) + assert.NotEmpty(t, encoded) + + decoded, err := DecodeV4(encoded) + require.NoError(t, err) + assert.Equal(t, c.ObjectType, decoded.ObjectType) + assert.Equal(t, c.ObjectID, decoded.ObjectID) + assert.Equal(t, c.RowID, decoded.RowID) + assert.Equal(t, c.SpanID, decoded.SpanID) + assert.Equal(t, c.RootSpanID, decoded.RootSpanID) +} + +func TestDecodeV4_MinimalPayload(t *testing.T) { + c := SpanComponents{ObjectType: SpanObjectTypeProjectLogs, ObjectID: "p1"} + encoded, err := EncodeV4(c) + require.NoError(t, err) + decoded, err := DecodeV4(encoded) + require.NoError(t, err) + assert.Equal(t, "p1", decoded.ObjectID) + assert.Equal(t, SpanObjectTypeProjectLogs, decoded.ObjectType) +} + +func TestEncodeDecodeV4_WithSpanIds(t *testing.T) { + c := SpanComponents{ + ObjectType: SpanObjectTypeExperiment, + ObjectID: "exp-uuid-here", + RowID: "1234567890abcdef", + SpanID: "1234567890abcdef", + RootSpanID: "a1b2c3d4e5f6789012345678abcdef01", + } + encoded, err := EncodeV4(c) + require.NoError(t, err) + decoded, err := DecodeV4(encoded) + require.NoError(t, err) + assert.Equal(t, c.SpanID, decoded.SpanID) + assert.Equal(t, c.RootSpanID, decoded.RootSpanID) + assert.Equal(t, SpanObjectTypeExperiment, decoded.ObjectType) +} + +func TestParentFromComponents_ProjectName(t *testing.T) { + c := SpanComponents{ObjectType: SpanObjectTypeProjectLogs, ObjectID: "my-project"} + p, err := ParentFromComponents(c) + require.NoError(t, err) + assert.Equal(t, ParentTypeProjectName, p.Type) + assert.Equal(t, "my-project", p.ID) +} + +func TestParentFromComponents_ProjectID(t *testing.T) { + c := SpanComponents{ObjectType: SpanObjectTypeProjectLogs, ObjectID: "550e8400-e29b-41d4-a716-446655440000"} + p, err := ParentFromComponents(c) + require.NoError(t, err) + assert.Equal(t, ParentTypeProjectID, p.Type) + assert.Equal(t, "550e8400-e29b-41d4-a716-446655440000", p.ID) +} + +func TestParentFromComponents_Experiment(t *testing.T) { + c := SpanComponents{ObjectType: SpanObjectTypeExperiment, ObjectID: "exp-123"} + p, err := ParentFromComponents(c) + require.NoError(t, err) + assert.Equal(t, ParentTypeExperimentID, p.Type) + assert.Equal(t, "exp-123", p.ID) +} diff --git a/trace/trace.go b/trace/trace.go index 84fe479..4994fcc 100644 --- a/trace/trace.go +++ b/trace/trace.go @@ -660,3 +660,88 @@ func getSpanURLData(span oteltrace.Span) (url, org string, parent Parent, err er org = attrs[orgAttrKey] return } + +// parentToSpanObjectType maps a Parent to the SpanComponents V4 object type. +func parentToSpanObjectType(p Parent) SpanObjectTypeV3 { + switch p.Type { + case ParentTypeExperimentID: + return SpanObjectTypeExperiment + case ParentTypeProjectName, ParentTypeProjectID: + return SpanObjectTypeProjectLogs + default: + return SpanObjectTypeProjectLogs + } +} + +// Export serializes the span into a string that can be passed to another process +// and used with ContextWithExportedSpan to create child spans there. The format +// is compatible with the Braintrust JS/Python span.export() output. +// +// Use this for distributed tracing: export the span in process A, pass the string +// to process B, then in B call ctx = trace.ContextWithExportedSpan(ctx, exported) +// and start spans with that context so they appear as children of the exported span. +// +// Export uses the span's Braintrust parent (project/experiment) and OTel trace/span +// IDs. The span must be recording and have Braintrust attributes (braintrust.parent, etc.) +// set by the Braintrust span processor. +func Export(span oteltrace.Span) (string, error) { + _, _, parent, err := getSpanURLData(span) + if err != nil { + return "", fmt.Errorf("get span data: %w", err) + } + spanContext := span.SpanContext() + if !spanContext.TraceID().IsValid() || !spanContext.SpanID().IsValid() { + return "", fmt.Errorf("span has invalid trace or span id") + } + traceID := spanContext.TraceID().String() + spanID := spanContext.SpanID().String() + + c := SpanComponents{ + ObjectType: parentToSpanObjectType(parent), + ObjectID: parent.ID, + RowID: spanID, + SpanID: spanID, + RootSpanID: traceID, + } + return EncodeV4(c) +} + +// ContextWithExportedSpan returns a context that carries the exported span as the +// remote parent. Spans started with this context (e.g. tracer.Start(ctx, "name")) +// will be children of that span and will be sent to the same Braintrust project/experiment. +// +// The exported string should be the result of Export(span) from this SDK or +// span.export() from the JS/Python SDK. +func ContextWithExportedSpan(ctx context.Context, exported string) (context.Context, error) { + components, err := DecodeV4(exported) + if err != nil { + return ctx, fmt.Errorf("decode exported span: %w", err) + } + if components.SpanID == "" || components.RootSpanID == "" { + return ctx, fmt.Errorf("exported span must contain span_id and root_span_id for use as parent") + } + + traceID, err := oteltrace.TraceIDFromHex(components.RootSpanID) + if err != nil { + return ctx, fmt.Errorf("invalid root_span_id (trace id): %w", err) + } + spanID, err := oteltrace.SpanIDFromHex(components.SpanID) + if err != nil { + return ctx, fmt.Errorf("invalid span_id: %w", err) + } + + remoteCtx := oteltrace.NewSpanContext(oteltrace.SpanContextConfig{ + TraceID: traceID, + SpanID: spanID, + TraceFlags: oteltrace.FlagsSampled, + }) + ctx = oteltrace.ContextWithRemoteSpanContext(ctx, remoteCtx) + + parent, err := ParentFromComponents(components) + if err != nil { + return ctx, err + } + ctx = SetParent(ctx, parent) + + return ctx, nil +} diff --git a/trace/trace_test.go b/trace/trace_test.go index e8091a1..ab92576 100644 --- a/trace/trace_test.go +++ b/trace/trace_test.go @@ -670,3 +670,86 @@ func TestPermalink_NoopSpan(t *testing.T) { noopSpan.End() } + +func TestExport(t *testing.T) { + assert := assert.New(t) + + tp := sdktrace.NewTracerProvider() + exporter := tracetest.NewInMemoryExporter() + + session := newTestSession() + cfg := Config{ + DefaultProjectID: "export-test-project", + Exporter: exporter, + Logger: logger.Discard(), + } + + err := AddSpanProcessor(tp, session, cfg) + assert.NoError(err) + + tracer := tp.Tracer("test") + _, span := tracer.Start(context.Background(), "parent-operation") + + exported, err := Export(span) + assert.NoError(err) + assert.NotEmpty(exported) + + components, err := DecodeV4(exported) + assert.NoError(err) + assert.Equal(SpanObjectTypeProjectLogs, components.ObjectType) + assert.Equal("export-test-project", components.ObjectID) + assert.Equal(span.SpanContext().SpanID().String(), components.SpanID) + assert.Equal(span.SpanContext().TraceID().String(), components.RootSpanID) + + span.End() +} + +func TestExport_NoopSpanFails(t *testing.T) { + noopTP := noop.NewTracerProvider() + noopTracer := noopTP.Tracer("test") + _, noopSpan := noopTracer.Start(context.Background(), "noop") + + _, err := Export(noopSpan) + assert.Error(t, err) + + noopSpan.End() +} + +func TestContextWithExportedSpan(t *testing.T) { + assert := assert.New(t) + + tp := sdktrace.NewTracerProvider() + exporter := tracetest.NewInMemoryExporter() + + session := newTestSession() + cfg := Config{ + DefaultProjectID: "ctx-export-project", + Exporter: exporter, + Logger: logger.Discard(), + } + + err := AddSpanProcessor(tp, session, cfg) + assert.NoError(err) + + tracer := tp.Tracer("test") + _, parentSpan := tracer.Start(context.Background(), "parent") + exported, err := Export(parentSpan) + assert.NoError(err) + parentSpan.End() + + components, err := DecodeV4(exported) + assert.NoError(err) + + childCtx, err := ContextWithExportedSpan(context.Background(), exported) + assert.NoError(err) + + _, childSpan := tracer.Start(childCtx, "child") + childSpan.End() + + // Child should be in same trace as the exported parent + assert.Equal(components.RootSpanID, childSpan.SpanContext().TraceID().String()) + + _ = tp.ForceFlush(context.Background()) + spans := exporter.GetSpans() + assert.GreaterOrEqual(len(spans), 1) +}