Support NoSample for span and its child spans

client_test.go

@@ -4,6 +4,7 @@
 package cozeloop
 
 import (
+	"context"
 	"testing"
 
 	. "github.com/smartystreets/goconvey/convey"
@@ -22,3 +23,28 @@ func TestNewClient(t *testing.T) {
 		So(client1, ShouldNotEqual, client3)
 	})
 }
+
+func TestNoSample(t *testing.T) {
+	Convey("Test NoSample spans", t, func() {
+		client, err := NewClient(WithWorkspaceID("test-workspace"), WithAPIToken("test-token"))
+		So(err, ShouldBeNil)
+		So(client, ShouldNotBeNil)
+
+		ctx := context.Background()
+
+		// 1. Create a normal span
+		normalCtx, normalSpan := client.StartSpan(ctx, "normal-span", "test-type")
+		So(normalSpan, ShouldNotBeNil)
+		So(normalSpan.GetSpanID(), ShouldNotBeEmpty)
+
+		// 2. Create a NoSample span using WithNoSample()
+		noSampleCtx, noSampleSpan := client.StartSpan(normalCtx, "nosample-span", "test-type", WithNoSample())
+		So(noSampleSpan, ShouldNotBeNil)
+		So(noSampleSpan.GetSpanID(), ShouldBeEmpty) // NoopSpan has empty SpanID
+
+		// 3. Create a child span from NoSample span - should also be NoopSpan
+		_, childSpan := client.StartSpan(noSampleCtx, "child-span", "test-type")
+		So(childSpan, ShouldNotBeNil)
+		So(childSpan.GetSpanID(), ShouldBeEmpty) // Child of NoopSpan is also NoopSpan
+	})
+}

internal/prompt/prompt_hub.go

@@ -9,6 +9,7 @@ import (
 	"io"
 	"time"
 
+	"github.com/coze-dev/cozeloop-go/internal/span"
 	"github.com/valyala/fasttemplate"
 
 	"github.com/coze-dev/cozeloop-go/entity"
@@ -60,7 +61,7 @@ func NewPromptProvider(httpClient *httpclient.Client, traceProvider *trace.Provi
 
 func (p *Provider) GetPrompt(ctx context.Context, param GetPromptParam, options GetPromptOptions) (prompt *entity.Prompt, err error) {
 	if p.config.PromptTrace && p.traceProvider != nil {
-		var promptHubSpan *trace.Span
+		var promptHubSpan span.Span
 		var spanErr error
 		ctx, promptHubSpan, spanErr = p.traceProvider.StartSpan(ctx, consts.TracePromptHubSpanName, tracespec.VPromptHubSpanType,
 			trace.StartSpanOptions{Scene: tracespec.VScenePromptHub})
@@ -135,7 +136,7 @@ func (p *Provider) PromptFormat(ctx context.Context, prompt *entity.Prompt, vari
 		return nil, nil
 	}
 	if p.config.PromptTrace && p.traceProvider != nil {
-		var promptTemplateSpan *trace.Span
+		var promptTemplateSpan span.Span
 		var spanErr error
 		ctx, promptTemplateSpan, spanErr = p.traceProvider.StartSpan(ctx, consts.TracePromptTemplateSpanName, tracespec.VPromptTemplateSpanType,
 			trace.StartSpanOptions{Scene: tracespec.VScenePromptTemplate})

internal/span/span.go

@@ -0,0 +1,140 @@
+// Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+// SPDX-License-Identifier: MIT
+
+package span
+
+import (
+	"context"
+	"time"
+
+	"github.com/coze-dev/cozeloop-go/entity"
+	"github.com/coze-dev/cozeloop-go/spec/tracespec"
+)
+
+// Span is the interface for span.
+type Span interface {
+	SpanContext
+	commonSpanSetter
+
+	// SetTags sets business custom tags.
+	SetTags(ctx context.Context, tagKVs map[string]interface{})
+
+	// SetBaggage sets tags and also passes these tags to other downstream spans (assuming
+	// the user uses ToHeader and FromHeader to handle header passing between services).
+	SetBaggage(ctx context.Context, baggageItems map[string]string)
+
+	// Finish The span will be reported only after an explicit call to Finish.
+	// Under the hood, it is actually placed in an asynchronous queue waiting to be reported.
+	Finish(ctx context.Context)
+
+	// GetStartTime returns the start time of the Span.
+	GetStartTime() time.Time
+
+	// ToHeader Convert the span to headers. Used for cross-process correlation.
+	ToHeader() (map[string]string, error)
+}
+
+// Set system-defined fields
+type commonSpanSetter interface {
+	// SetInput key: `input`
+	// Input information. The input will be serialized into a JSON string.
+	// You can find recommended specification in https://github.com/coze-dev/cozeloop-go/tree/main/spec/tracespec
+	// Or you can use any struct you like.
+	SetInput(ctx context.Context, input interface{})
+
+	// SetOutput key: `output`
+	// Output information. The output will be serialized into a JSON string.
+	// You can find recommended specification in https://github.com/coze-dev/cozeloop-go/tree/main/spec/tracespec
+	// Or you can use any struct you like.
+	SetOutput(ctx context.Context, output interface{})
+
+	// SetError key: `error`
+	// Set error message.
+	SetError(ctx context.Context, err error)
+
+	// SetStatusCode key: `status_code`
+	// Set status code. A non-zero code is considered an exception.
+	SetStatusCode(ctx context.Context, code int)
+
+	// SetUserID key: `user_id`
+	// Set user id.
+	SetUserID(ctx context.Context, userID string)
+	SetUserIDBaggage(ctx context.Context, userID string)
+
+	// SetMessageID key: `message_id`
+	// Set message id.
+	SetMessageID(ctx context.Context, messageID string)
+	SetMessageIDBaggage(ctx context.Context, messageID string)
+
+	// SetThreadID key: `thread_id`
+	// Set thread id.
+	// thread_id is used to correlate multiple requests and is passed in by the business.
+	SetThreadID(ctx context.Context, threadID string)
+	SetThreadIDBaggage(ctx context.Context, threadID string)
+
+	// SetPrompt key: `prompt
+	// Associated with PromptKey and PromptVersion, it will write two tags: prompt_key and prompt_version.
+	// SetPrompt is used to set the PromptKey and PromptVersion to tag.
+	SetPrompt(ctx context.Context, prompt entity.Prompt)
+
+	// SetModelProvider key: `model_provider`
+	// The provider of the LLM, such as OpenAI, etc.
+	SetModelProvider(ctx context.Context, modelProvider string)
+
+	// SetModelName key: `model_name`
+	// The Name of the LLM model, such as: gpt-4-1106-preview.
+	SetModelName(ctx context.Context, modelName string)
+
+	// SetModelCallOptions key: `call_options`
+	// The call options of the LLM, such as temperature, max_tokens, etc.
+	// The recommended standard format is CallOption of spec package
+	SetModelCallOptions(ctx context.Context, callOptions interface{})
+
+	// SetInputTokens key: `input_tokens`
+	// The usage of input tokens. When the value of input_tokens is set,
+	// It will be automatically summed with output_tokens to calculate the tokens tag.
+	SetInputTokens(ctx context.Context, inputTokens int)
+
+	// SetOutputTokens key: `output_tokens`
+	// The usage of output tokens. When the value of output_tokens is set,
+	// It will be automatically summed with input_tokens to calculate the tokens tag.
+	SetOutputTokens(ctx context.Context, outputTokens int)
+
+	// SetStartTimeFirstResp key: `start_time_first_resp`
+	// Timestamp of the first packet return from LLM, unit: microseconds.
+	// When `start_time_first_resp` is set, a tag named `latency_first_resp` calculated
+	// based on the span's StartTime will be added, meaning the latency for the first packet.
+	SetStartTimeFirstResp(ctx context.Context, startTimeFirstResp int64)
+
+	// SetRuntime key: `runtime`
+	// The runtime of the LLM, such as language, library, scene, etc.
+	// The recommended standard format is Runtime of spec package
+	SetRuntime(ctx context.Context, runtime tracespec.Runtime)
+
+	// SetServiceName
+	// set the custom service name, identify different services.
+	SetServiceName(ctx context.Context, serviceName string)
+
+	// SetLogID
+	// set the custom log id, identify different query.
+	SetLogID(ctx context.Context, logID string)
+
+	// SetFinishTime
+	// Default is time.Now() when span Finish(). DO NOT set unless you do not use default time.
+	SetFinishTime(finishTime time.Time)
+
+	// SetSystemTags
+	// set the system tags. DO NOT set unless you know what you are doing.
+	SetSystemTags(ctx context.Context, systemTags map[string]interface{})
+
+	// SetDeploymentEnv
+	// set the deployment env, identify custom env.
+	SetDeploymentEnv(ctx context.Context, deploymentEnv string)
+}
+
+// SpanContext is the interface for span Baggage transfer.
+type SpanContext interface {
+	GetSpanID() string
+	GetTraceID() string
+	GetBaggage() map[string]string
+}

internal/trace/exporter_test.go

@@ -8,13 +8,14 @@ import (
 	"testing"
 
 	. "github.com/bytedance/mockey"
+	"github.com/coze-dev/cozeloop-go/entity"
 	"github.com/coze-dev/cozeloop-go/internal/httpclient"
 	. "github.com/smartystreets/goconvey/convey"
 )
 
 func Test_ExportSpans(t *testing.T) {
 	ctx := context.Background()
-	spans := []*UploadSpan{{}, {}}
+	spans := []*entity.UploadSpan{{}, {}}
 
 	PatchConvey("Test transferToUploadSpanAndFile failed", t, func() {
 		Mock((*httpclient.Client).Post).Return(nil).Build()

internal/trace/noop_span.go

@@ -8,11 +8,14 @@ import (
 	"time"
 
 	"github.com/coze-dev/cozeloop-go/entity"
+	"github.com/coze-dev/cozeloop-go/internal/span"
 	"github.com/coze-dev/cozeloop-go/spec/tracespec"
 )
 
 var DefaultNoopSpan = &noopSpan{}
 
+var _ span.Span = noopSpan{}
+
 type noopSpan struct{}
 
 // implement of commonSpanSetter

internal/trace/queue_manager_test.go

@@ -15,7 +15,7 @@ import (
 func Test_GetBatchSpanProcessor(t *testing.T) {
 	ctx := context.Background()
 	httpClient := &httpclient.Client{}
-	spanQM := NewBatchSpanProcessor(nil, httpClient, nil, nil)
+	spanQM := NewBatchSpanProcessor(nil, httpClient, nil, nil, nil)
 
 	PatchConvey("Test GetBatchSpanProcessor", t, func() {
 		PatchConvey("Test with valid inputs", func() {

internal/trace/span.go

@@ -19,6 +19,7 @@ import (
 	"github.com/coze-dev/cozeloop-go/internal"
 	"github.com/coze-dev/cozeloop-go/internal/consts"
 	"github.com/coze-dev/cozeloop-go/internal/logger"
+	"github.com/coze-dev/cozeloop-go/internal/span"
 	"github.com/coze-dev/cozeloop-go/internal/util"
 	"github.com/coze-dev/cozeloop-go/spec/tracespec"
 )
@@ -28,6 +29,8 @@ const (
 	spanFinished   = 1
 )
 
+var _ span.SpanContext = (*SpanContext)(nil)
+
 type SpanContext struct {
 	SpanID  string
 	TraceID string
@@ -46,6 +49,8 @@ func (s *SpanContext) GetBaggage() map[string]string {
 	return s.Baggage
 }
 
+var _ span.Span = (*Span)(nil)
+
 type Span struct {
 	// span context param
 	SpanContext

internal/trace/span_test.go

@@ -129,7 +129,7 @@ func Test_Finish(t *testing.T) {
 	httpClient := httpclient.NewClient("", nil, nil, nil)
 	s := &Span{
 		isFinished:    0,
-		spanProcessor: NewBatchSpanProcessor(nil, httpClient, nil, nil),
+		spanProcessor: NewBatchSpanProcessor(nil, httpClient, nil, nil, nil),
 		lock:          sync.RWMutex{},
 		TagMap:        make(map[string]interface{}),
 	}

internal/trace/trace.go

@@ -11,6 +11,7 @@ import (
 	"github.com/coze-dev/cozeloop-go/internal/consts"
 	"github.com/coze-dev/cozeloop-go/internal/httpclient"
 	"github.com/coze-dev/cozeloop-go/internal/logger"
+	"github.com/coze-dev/cozeloop-go/internal/span"
 	"github.com/coze-dev/cozeloop-go/internal/util"
 	"github.com/coze-dev/cozeloop-go/spec/tracespec"
 )
@@ -41,6 +42,7 @@ type StartSpanOptions struct {
 	StartNewTrace bool
 	Scene         string
 	WorkspaceID   string
+	NoSample      bool // whether to disable sampling for this span and its child spans
 }
 
 type loopSpanKey struct{}
@@ -71,7 +73,16 @@ func (t *Provider) GetOpts() *Options {
 	return t.opt
 }
 
-func (t *Provider) StartSpan(ctx context.Context, name, spanType string, opts StartSpanOptions) (context.Context, *Span, error) {
+func (t *Provider) StartSpan(ctx context.Context, name, spanType string, opts StartSpanOptions) (context.Context, span.Span, error) {
+	if opts.NoSample {
+		ctx = context.WithValue(ctx, loopSpanKey{}, DefaultNoopSpan)
+		return ctx, DefaultNoopSpan, nil
+	}
+
+	if !t.shouldSample(ctx) {
+		return ctx, DefaultNoopSpan, nil
+	}
+
 	// 0. check param
 	if name == "" {
 		name = "unknown"
@@ -125,6 +136,11 @@ func (t *Provider) GetSpanFromHeader(ctx context.Context, header map[string]stri
 	return FromHeader(ctx, header)
 }
 
+func (t *Provider) shouldSample(ctx context.Context) bool {
+	parentSpan := ctx.Value(loopSpanKey{})
+	return parentSpan != DefaultNoopSpan
+}
+
 func (t *Provider) startSpan(ctx context.Context, spanName string, spanType string, options StartSpanOptions) *Span {
 	// 1. pack base data
 	// get parentID from opt first, or set it to "0".

internal/trace/trace_test.go

@@ -76,3 +76,40 @@ func Test_GetSpanFromHeader(t *testing.T) {
 		So(actual, ShouldEqual, expectedSpan)
 	})
 }
+
+func Test_NoSample(t *testing.T) {
+	PatchConvey("Test NoSample spans", t, func() {
+		provider := &Provider{
+			httpClient: &httpclient.Client{},
+			opt: &Options{
+				WorkspaceID:      "workspace-id",
+				UltraLargeReport: true,
+			},
+		}
+		ctx := context.Background()
+
+		// 1. Create a normal span
+		normalCtx, normalSpan, err := provider.StartSpan(ctx, "normal-span", "test-type", StartSpanOptions{
+			StartTime: time.Now(),
+		})
+		So(err, ShouldBeNil)
+		So(normalSpan, ShouldNotEqual, DefaultNoopSpan)
+		So(normalSpan, ShouldNotBeNil)
+
+		// 2. Create a NoSample span
+		noSampleCtx, noSampleSpan, err := provider.StartSpan(normalCtx, "nosample-span", "test-type", StartSpanOptions{
+			StartTime: time.Now(),
+			NoSample:  true,
+		})
+		So(err, ShouldBeNil)
+		So(noSampleSpan, ShouldEqual, DefaultNoopSpan)
+
+		// 3. Create a child span from NoSample span - should also be NoopSpan
+		childCtx, childSpan, err := provider.StartSpan(noSampleCtx, "child-span", "test-type", StartSpanOptions{
+			StartTime: time.Now(),
+		})
+		So(err, ShouldBeNil)
+		So(childSpan, ShouldEqual, DefaultNoopSpan)
+		So(childCtx, ShouldNotBeNil)
+	})
+}