From ca0158994b2f76e51b2dca2c1db35f52e121ecdd Mon Sep 17 00:00:00 2001
From: Parship Chowdhury <parshipchowdhury@gmail.com>
Date: Wed, 31 Dec 2025 19:25:49 +0000
Subject: [PATCH 01/22] implement max_trace_size parameter for v2 query service

Signed-off-by: Parship Chowdhury <parshipchowdhury@gmail.com>
---
 .../extension/jaegerquery/internal/flags.go   |   3 +
 .../internal/querysvc/v2/querysvc/service.go  |  34 +++-
 .../querysvc/v2/querysvc/service_test.go      | 159 ++++++++++++++++++
 .../internal/extension/jaegerquery/server.go  |   1 +
 internal/jptrace/aggregator.go                | 131 ++++++++++++++-
 internal/jptrace/aggregator_test.go           |  76 +++++++++
 6 files changed, 398 insertions(+), 6 deletions(-)

diff --git a/cmd/jaeger/internal/extension/jaegerquery/internal/flags.go b/cmd/jaeger/internal/extension/jaegerquery/internal/flags.go
index 502738bc7fd..bb65f9c0d68 100644
--- a/cmd/jaeger/internal/extension/jaegerquery/internal/flags.go
+++ b/cmd/jaeger/internal/extension/jaegerquery/internal/flags.go
@@ -36,6 +36,9 @@ type QueryOptions struct {
 	Tenancy tenancy.Options `mapstructure:"multi_tenancy"`
 	// MaxClockSkewAdjust is the maximum duration by which jaeger-query will adjust a span.
 	MaxClockSkewAdjust time.Duration `mapstructure:"max_clock_skew_adjust"  valid:"optional"`
+	// MaxTraceSize is the max no. of spans allowed per trace.
+	// If a trace has more spans than this, it will be truncated and a warning will be added
+	MaxTraceSize int `mapstructure:"max_trace_size" valid:"optional"`
 	// EnableTracing determines whether traces will be emitted by jaeger-query.
 	EnableTracing bool `mapstructure:"enable_tracing"`
 	// HTTP holds the HTTP configuration that the query service uses to serve requests.
diff --git a/cmd/jaeger/internal/extension/jaegerquery/internal/querysvc/v2/querysvc/service.go b/cmd/jaeger/internal/extension/jaegerquery/internal/querysvc/v2/querysvc/service.go
index 58b70aa259f..313987a75ca 100644
--- a/cmd/jaeger/internal/extension/jaegerquery/internal/querysvc/v2/querysvc/service.go
+++ b/cmd/jaeger/internal/extension/jaegerquery/internal/querysvc/v2/querysvc/service.go
@@ -6,6 +6,7 @@ package querysvc
 import (
 	"context"
 	"errors"
+	"fmt"
 	"iter"
 	"time"
 
@@ -29,6 +30,9 @@ type QueryServiceOptions struct {
 	ArchiveTraceWriter tracestore.Writer
 	// MaxClockSkewAdjust is the maximum duration by which to adjust a span.
 	MaxClockSkewAdjust time.Duration
+	// MaxTraceSize is the max no. of spans allowed per trace.
+	// If a trace has more spans than this, it will be truncated and a warning will be added
+	MaxTraceSize int
 }
 
 // StorageCapabilities is a feature flag for query service
@@ -200,10 +204,38 @@ func (qs QueryService) receiveTraces(
 	if rawTraces {
 		seq(processTraces)
 	} else {
-		jptrace.AggregateTraces(seq)(func(trace ptrace.Traces, err error) bool {
+		jptrace.AggregateTracesWithLimit(seq, qs.options.MaxTraceSize)(func(trace ptrace.Traces, err error) bool {
+			// Add warning if trace was truncated
+			if err == nil && qs.options.MaxTraceSize > 0 && jptrace.IsTraceTruncated(trace) {
+				qs.addTruncationWarning(trace)
+			}
 			return processTraces([]ptrace.Traces{trace}, err)
 		})
 	}
 
 	return foundTraceIDs, proceed
 }
+
+// add a warning to the first span of the trace
+func (qs QueryService) addTruncationWarning(trace ptrace.Traces) {
+	resources := trace.ResourceSpans()
+	if resources.Len() == 0 {
+		return
+	}
+
+	scopes := resources.At(0).ScopeSpans()
+	if scopes.Len() == 0 {
+		return
+	}
+
+	spans := scopes.At(0).Spans()
+	if spans.Len() == 0 {
+		return
+	}
+
+	firstSpan := spans.At(0)
+	firstSpan.Attributes().Remove("@jaeger@truncated")
+	jptrace.AddWarnings(firstSpan,
+		fmt.Sprintf("trace has more than %d spans, showing first %d spans only",
+			qs.options.MaxTraceSize, qs.options.MaxTraceSize))
+}
diff --git a/cmd/jaeger/internal/extension/jaegerquery/internal/querysvc/v2/querysvc/service_test.go b/cmd/jaeger/internal/extension/jaegerquery/internal/querysvc/v2/querysvc/service_test.go
index dfd096e197a..34779b3f977 100644
--- a/cmd/jaeger/internal/extension/jaegerquery/internal/querysvc/v2/querysvc/service_test.go
+++ b/cmd/jaeger/internal/extension/jaegerquery/internal/querysvc/v2/querysvc/service_test.go
@@ -623,3 +623,162 @@ func TestGetCapabilities(t *testing.T) {
 		})
 	}
 }
+
+func TestMaxTraceSize_UnderLimit(t *testing.T) {
+	// 3 spans
+	trace := ptrace.NewTraces()
+	resources := trace.ResourceSpans().AppendEmpty()
+	scopes := resources.ScopeSpans().AppendEmpty()
+	for i := 0; i < 3; i++ {
+		span := scopes.Spans().AppendEmpty()
+		span.SetTraceID(testTraceID)
+		span.SetSpanID(pcommon.SpanID([8]byte{byte(i + 1)}))
+		span.SetName(fmt.Sprintf("span-%d", i))
+	}
+
+	responseIter := iter.Seq2[[]ptrace.Traces, error](func(yield func([]ptrace.Traces, error) bool) {
+		yield([]ptrace.Traces{trace}, nil)
+	})
+
+	traceReader := &tracestoremocks.Reader{}
+	dependencyStorage := &depstoremocks.Reader{}
+	options := QueryServiceOptions{
+		MaxTraceSize: 5, // limit is 5, but trace has only 3 spans
+	}
+	tqs := &testQueryService{}
+	tqs.queryService = NewQueryService(traceReader, dependencyStorage, options)
+
+	params := GetTraceParams{
+		TraceIDs: []tracestore.GetTraceParams{{TraceID: testTraceID}},
+	}
+	traceReader.On("GetTraces", mock.Anything, params.TraceIDs).
+		Return(responseIter).Once()
+
+	getTracesIter := tqs.queryService.GetTraces(context.Background(), params)
+	gotTraces, err := jiter.FlattenWithErrors(getTracesIter)
+	require.NoError(t, err)
+	require.Len(t, gotTraces, 1)
+
+	gotSpans := gotTraces[0].ResourceSpans().At(0).ScopeSpans().At(0).Spans()
+	require.Equal(t, 3, gotSpans.Len())
+
+	// no warning should be present
+	firstSpan := gotSpans.At(0)
+	_, hasWarning := firstSpan.Attributes().Get("@jaeger@warnings")
+	require.False(t, hasWarning)
+}
+
+func TestMaxTraceSize_OverLimit(t *testing.T) {
+	// 5 spans split across 2 batches
+	trace1 := ptrace.NewTraces()
+	resources1 := trace1.ResourceSpans().AppendEmpty()
+	scopes1 := resources1.ScopeSpans().AppendEmpty()
+	for i := 0; i < 3; i++ {
+		span := scopes1.Spans().AppendEmpty()
+		span.SetTraceID(testTraceID)
+		span.SetSpanID(pcommon.SpanID([8]byte{byte(i + 1)}))
+		span.SetName(fmt.Sprintf("span-%d", i))
+	}
+
+	trace2 := ptrace.NewTraces()
+	resources2 := trace2.ResourceSpans().AppendEmpty()
+	scopes2 := resources2.ScopeSpans().AppendEmpty()
+	for i := 3; i < 5; i++ {
+		span := scopes2.Spans().AppendEmpty()
+		span.SetTraceID(testTraceID)
+		span.SetSpanID(pcommon.SpanID([8]byte{byte(i + 1)}))
+		span.SetName(fmt.Sprintf("span-%d", i))
+	}
+
+	responseIter := iter.Seq2[[]ptrace.Traces, error](func(yield func([]ptrace.Traces, error) bool) {
+		if !yield([]ptrace.Traces{trace1}, nil) {
+			return
+		}
+		yield([]ptrace.Traces{trace2}, nil)
+	})
+
+	traceReader := &tracestoremocks.Reader{}
+	dependencyStorage := &depstoremocks.Reader{}
+	options := QueryServiceOptions{
+		MaxTraceSize: 3, // Limit is 3, but trace has 5 spans total
+	}
+	tqs := &testQueryService{}
+	tqs.queryService = NewQueryService(traceReader, dependencyStorage, options)
+
+	params := GetTraceParams{
+		TraceIDs: []tracestore.GetTraceParams{{TraceID: testTraceID}},
+	}
+	traceReader.On("GetTraces", mock.Anything, params.TraceIDs).
+		Return(responseIter).Once()
+
+	getTracesIter := tqs.queryService.GetTraces(context.Background(), params)
+	gotTraces, err := jiter.FlattenWithErrors(getTracesIter)
+	require.NoError(t, err)
+	require.Len(t, gotTraces, 1)
+
+	// count total spans in the result
+	totalSpans := 0
+	resources := gotTraces[0].ResourceSpans()
+	for i := 0; i < resources.Len(); i++ {
+		scopes := resources.At(i).ScopeSpans()
+		for j := 0; j < scopes.Len(); j++ {
+			totalSpans += scopes.At(j).Spans().Len()
+		}
+	}
+
+	// Only 3 spans should be present(the limit)
+	require.Equal(t, 3, totalSpans)
+
+	// there should be a warning for the first span
+	firstSpan := resources.At(0).ScopeSpans().At(0).Spans().At(0)
+	warningsAttr, hasWarning := firstSpan.Attributes().Get("@jaeger@warnings")
+	require.True(t, hasWarning)
+	require.Equal(t, pcommon.ValueTypeSlice, warningsAttr.Type())
+	warnings := warningsAttr.Slice()
+	require.Positive(t, warnings.Len())
+	require.Contains(t, warnings.At(warnings.Len()-1).Str(), "trace has more than 3 spans")
+}
+
+func TestMaxTraceSize_ExactlyAtLimit(t *testing.T) {
+	// 3 spans
+	trace := ptrace.NewTraces()
+	resources := trace.ResourceSpans().AppendEmpty()
+	scopes := resources.ScopeSpans().AppendEmpty()
+	for i := 0; i < 3; i++ {
+		span := scopes.Spans().AppendEmpty()
+		span.SetTraceID(testTraceID)
+		span.SetSpanID(pcommon.SpanID([8]byte{byte(i + 1)}))
+		span.SetName(fmt.Sprintf("span-%d", i))
+	}
+
+	responseIter := iter.Seq2[[]ptrace.Traces, error](func(yield func([]ptrace.Traces, error) bool) {
+		yield([]ptrace.Traces{trace}, nil)
+	})
+
+	traceReader := &tracestoremocks.Reader{}
+	dependencyStorage := &depstoremocks.Reader{}
+	options := QueryServiceOptions{
+		MaxTraceSize: 3, // Limit is exactly 3, trace has 3 spans
+	}
+	tqs := &testQueryService{}
+	tqs.queryService = NewQueryService(traceReader, dependencyStorage, options)
+
+	params := GetTraceParams{
+		TraceIDs: []tracestore.GetTraceParams{{TraceID: testTraceID}},
+	}
+	traceReader.On("GetTraces", mock.Anything, params.TraceIDs).
+		Return(responseIter).Once()
+
+	getTracesIter := tqs.queryService.GetTraces(context.Background(), params)
+	gotTraces, err := jiter.FlattenWithErrors(getTracesIter)
+	require.NoError(t, err)
+	require.Len(t, gotTraces, 1)
+
+	gotSpans := gotTraces[0].ResourceSpans().At(0).ScopeSpans().At(0).Spans()
+	// all 3 spans should be present
+	require.Equal(t, 3, gotSpans.Len())
+
+	firstSpan := gotSpans.At(0)
+	_, hasWarning := firstSpan.Attributes().Get("@jaeger@warnings")
+	require.False(t, hasWarning)
+}
diff --git a/cmd/jaeger/internal/extension/jaegerquery/server.go b/cmd/jaeger/internal/extension/jaegerquery/server.go
index 22bbee4abe8..60d2b777bc9 100644
--- a/cmd/jaeger/internal/extension/jaegerquery/server.go
+++ b/cmd/jaeger/internal/extension/jaegerquery/server.go
@@ -106,6 +106,7 @@ func (s *server) Start(ctx context.Context, host component.Host) error {
 	}
 	v2opts := v2querysvc.QueryServiceOptions{
 		MaxClockSkewAdjust: s.config.MaxClockSkewAdjust,
+		MaxTraceSize:       s.config.MaxTraceSize,
 	}
 	if err := s.addArchiveStorage(&opts, &v2opts, host); err != nil {
 		return err
diff --git a/internal/jptrace/aggregator.go b/internal/jptrace/aggregator.go
index fa8cde63646..74bb4eb6332 100644
--- a/internal/jptrace/aggregator.go
+++ b/internal/jptrace/aggregator.go
@@ -10,13 +10,21 @@ import (
 	"go.opentelemetry.io/collector/pdata/ptrace"
 )
 
+const truncationMarkerKey = "@jaeger@truncated"
+
 // AggregateTraces aggregates a sequence of trace batches into individual traces.
 //
 // The `tracesSeq` input must adhere to the chunking requirements of tracestore.Reader.GetTraces.
 func AggregateTraces(tracesSeq iter.Seq2[[]ptrace.Traces, error]) iter.Seq2[ptrace.Traces, error] {
+	return AggregateTracesWithLimit(tracesSeq, 0)
+}
+
+func AggregateTracesWithLimit(tracesSeq iter.Seq2[[]ptrace.Traces, error], maxSize int) iter.Seq2[ptrace.Traces, error] {
 	return func(yield func(trace ptrace.Traces, err error) bool) {
 		currentTrace := ptrace.NewTraces()
 		currentTraceID := pcommon.NewTraceIDEmpty()
+		spanCount := 0
+		skipCurrentTrace := false
 
 		tracesSeq(func(traces []ptrace.Traces, err error) bool {
 			if err != nil {
@@ -27,7 +35,13 @@ func AggregateTraces(tracesSeq iter.Seq2[[]ptrace.Traces, error]) iter.Seq2[ptra
 				resources := trace.ResourceSpans()
 				traceID := resources.At(0).ScopeSpans().At(0).Spans().At(0).TraceID()
 				if currentTraceID == traceID {
-					mergeTraces(trace, currentTrace)
+					if !skipCurrentTrace {
+						truncated := mergeTraces(trace, currentTrace, maxSize, &spanCount)
+						if truncated {
+							markTraceTruncated(currentTrace)
+							skipCurrentTrace = true
+						}
+					}
 				} else {
 					if currentTrace.ResourceSpans().Len() > 0 {
 						if !yield(currentTrace, nil) {
@@ -36,6 +50,15 @@ func AggregateTraces(tracesSeq iter.Seq2[[]ptrace.Traces, error]) iter.Seq2[ptra
 					}
 					currentTrace = trace
 					currentTraceID = traceID
+					spanCount = countSpans(trace)
+					skipCurrentTrace = false
+					if maxSize > 0 && spanCount > maxSize {
+						currentTrace = ptrace.NewTraces()
+						copySpansUpToLimit(trace, currentTrace, maxSize)
+						spanCount = maxSize
+						markTraceTruncated(currentTrace)
+						skipCurrentTrace = true
+					}
 				}
 			}
 			return true
@@ -46,10 +69,108 @@ func AggregateTraces(tracesSeq iter.Seq2[[]ptrace.Traces, error]) iter.Seq2[ptra
 	}
 }
 
-func mergeTraces(src, dest ptrace.Traces) {
-	resources := src.ResourceSpans()
+func mergeTraces(src, dest ptrace.Traces, maxSize int, spanCount *int) bool {
+	if maxSize <= 0 {
+		// No limit, merge all
+		resources := src.ResourceSpans()
+		for i := 0; i < resources.Len(); i++ {
+			resource := resources.At(i)
+			resource.CopyTo(dest.ResourceSpans().AppendEmpty())
+		}
+		return false
+	}
+
+	// with limit
+	incomingCount := countSpans(src)
+	if *spanCount+incomingCount <= maxSize {
+		resources := src.ResourceSpans()
+		for i := 0; i < resources.Len(); i++ {
+			resource := resources.At(i)
+			resource.CopyTo(dest.ResourceSpans().AppendEmpty())
+		}
+		*spanCount += incomingCount
+		return false
+	}
+
+	// partial copy
+	remaining := maxSize - *spanCount
+	if remaining > 0 {
+		copySpansUpToLimit(src, dest, remaining)
+		*spanCount = maxSize
+	}
+	return true
+}
+
+func countSpans(trace ptrace.Traces) int {
+	count := 0
+	resources := trace.ResourceSpans()
 	for i := 0; i < resources.Len(); i++ {
-		resource := resources.At(i)
-		resource.CopyTo(dest.ResourceSpans().AppendEmpty())
+		scopes := resources.At(i).ScopeSpans()
+		for j := 0; j < scopes.Len(); j++ {
+			count += scopes.At(j).Spans().Len()
+		}
+	}
+	return count
+}
+
+func copySpansUpToLimit(src, dest ptrace.Traces, limit int) {
+	copied := 0
+	resources := src.ResourceSpans()
+
+	for i := 0; i < resources.Len() && copied < limit; i++ {
+		srcResource := resources.At(i)
+		destResource := dest.ResourceSpans().AppendEmpty()
+		srcResource.Resource().CopyTo(destResource.Resource())
+		destResource.SetSchemaUrl(srcResource.SchemaUrl())
+
+		scopes := srcResource.ScopeSpans()
+		for j := 0; j < scopes.Len() && copied < limit; j++ {
+			srcScope := scopes.At(j)
+			destScope := destResource.ScopeSpans().AppendEmpty()
+			srcScope.Scope().CopyTo(destScope.Scope())
+			destScope.SetSchemaUrl(srcScope.SchemaUrl())
+
+			spans := srcScope.Spans()
+			for k := 0; k < spans.Len() && copied < limit; k++ {
+				spans.At(k).CopyTo(destScope.Spans().AppendEmpty())
+				copied++
+			}
+		}
+	}
+}
+
+func markTraceTruncated(trace ptrace.Traces) {
+	resources := trace.ResourceSpans()
+	if resources.Len() == 0 {
+		return
+	}
+	scopes := resources.At(0).ScopeSpans()
+	if scopes.Len() == 0 {
+		return
+	}
+	spans := scopes.At(0).Spans()
+	if spans.Len() == 0 {
+		return
+	}
+	firstSpan := spans.At(0)
+	firstSpan.Attributes().PutBool(truncationMarkerKey, true)
+}
+
+// check if a trace has marked as truncated
+func IsTraceTruncated(trace ptrace.Traces) bool {
+	resources := trace.ResourceSpans()
+	if resources.Len() == 0 {
+		return false
+	}
+	scopes := resources.At(0).ScopeSpans()
+	if scopes.Len() == 0 {
+		return false
+	}
+	spans := scopes.At(0).Spans()
+	if spans.Len() == 0 {
+		return false
 	}
+	firstSpan := spans.At(0)
+	val, exists := firstSpan.Attributes().Get(truncationMarkerKey)
+	return exists && val.Bool()
 }
diff --git a/internal/jptrace/aggregator_test.go b/internal/jptrace/aggregator_test.go
index abeeaaa0845..620794b9433 100644
--- a/internal/jptrace/aggregator_test.go
+++ b/internal/jptrace/aggregator_test.go
@@ -134,3 +134,79 @@ func TestAggregateTraces_RespectsEarlyReturn(t *testing.T) {
 
 	require.Equal(t, trace1, lastResult)
 }
+
+func TestAggregateTracesWithLimit(t *testing.T) {
+	createTrace := func(traceID byte, spanCount int) ptrace.Traces {
+		trace := ptrace.NewTraces()
+		spans := trace.ResourceSpans().AppendEmpty().ScopeSpans().AppendEmpty().Spans()
+		for i := 0; i < spanCount; i++ {
+			span := spans.AppendEmpty()
+			span.SetTraceID(pcommon.TraceID([16]byte{traceID}))
+		}
+		return trace
+	}
+
+	tests := []struct {
+		name           string
+		maxSize        int
+		inputSpans     int
+		expectedSpans  int
+		expectTruncate bool
+	}{
+		{"no_limit", 0, 5, 5, false},
+		{"under_limit", 10, 5, 5, false},
+		{"over_limit", 3, 5, 3, true},
+		{"exact_limit", 5, 5, 5, false},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			tracesSeq := func(yield func([]ptrace.Traces, error) bool) {
+				yield([]ptrace.Traces{createTrace(1, tt.inputSpans)}, nil)
+			}
+
+			var result []ptrace.Traces
+			AggregateTracesWithLimit(tracesSeq, tt.maxSize)(func(trace ptrace.Traces, _ error) bool {
+				result = append(result, trace)
+				return true
+			})
+
+			require.Len(t, result, 1)
+			assert.Equal(t, tt.expectedSpans, countSpans(result[0]))
+			assert.Equal(t, tt.expectTruncate, IsTraceTruncated(result[0]))
+		})
+	}
+}
+
+func TestCountSpans(t *testing.T) {
+	trace := ptrace.NewTraces()
+	r1 := trace.ResourceSpans().AppendEmpty()
+	r1.ScopeSpans().AppendEmpty().Spans().AppendEmpty()
+	r1.ScopeSpans().AppendEmpty().Spans().AppendEmpty()
+	r2 := trace.ResourceSpans().AppendEmpty()
+	r2.ScopeSpans().AppendEmpty().Spans().AppendEmpty()
+
+	assert.Equal(t, 3, countSpans(trace))
+}
+
+func TestCopySpansUpToLimit(t *testing.T) {
+	src := ptrace.NewTraces()
+	spans := src.ResourceSpans().AppendEmpty().ScopeSpans().AppendEmpty().Spans()
+	for i := 0; i < 5; i++ {
+		spans.AppendEmpty().SetName("span")
+	}
+
+	dest := ptrace.NewTraces()
+	copySpansUpToLimit(src, dest, 3)
+
+	assert.Equal(t, 3, countSpans(dest))
+}
+
+func TestMarkAndCheckTruncated(t *testing.T) {
+	trace := ptrace.NewTraces()
+	trace.ResourceSpans().AppendEmpty().ScopeSpans().AppendEmpty().Spans().AppendEmpty()
+
+	assert.False(t, IsTraceTruncated(trace))
+	markTraceTruncated(trace)
+	assert.True(t, IsTraceTruncated(trace))
+}

From 48362342d57244226b836c22d066f2d08478d76f Mon Sep 17 00:00:00 2001
From: Parship Chowdhury <parshipchowdhury@gmail.com>
Date: Thu, 1 Jan 2026 10:10:43 +0530
Subject: [PATCH 02/22] Update internal/jptrace/aggregator.go

Co-authored-by: Yuri Shkuro <yurishkuro@users.noreply.github.com>
Signed-off-by: Parship Chowdhury <parshipchowdhury@gmail.com>
---
 internal/jptrace/aggregator.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/internal/jptrace/aggregator.go b/internal/jptrace/aggregator.go
index 74bb4eb6332..c0af2ebd4bd 100644
--- a/internal/jptrace/aggregator.go
+++ b/internal/jptrace/aggregator.go
@@ -36,7 +36,7 @@ func AggregateTracesWithLimit(tracesSeq iter.Seq2[[]ptrace.Traces, error], maxSi
 				traceID := resources.At(0).ScopeSpans().At(0).Spans().At(0).TraceID()
 				if currentTraceID == traceID {
 					if !skipCurrentTrace {
-						truncated := mergeTraces(trace, currentTrace, maxSize, &spanCount)
+						spanCount = mergeTraces(trace, currentTrace, maxSize, spanCount)
 						if truncated {
 							markTraceTruncated(currentTrace)
 							skipCurrentTrace = true

From 74331284ff5f8cba971472dc5e3d93842f60220e Mon Sep 17 00:00:00 2001
From: Parship Chowdhury <parshipchowdhury@gmail.com>
Date: Thu, 1 Jan 2026 05:22:53 +0000
Subject: [PATCH 03/22] 1. remove countSpans and used trace.spanCount() 2.
 remove same code at aggregator.go 3. remove skipCurrentTrace and do an early
 exit from mergeTraces if the current count is already at max 4. remove
 redundant checks at aggregator.go

Signed-off-by: Parship Chowdhury <parshipchowdhury@gmail.com>
---
 internal/jptrace/aggregator.go      | 82 ++++++++---------------------
 internal/jptrace/aggregator_test.go | 15 +-----
 2 files changed, 24 insertions(+), 73 deletions(-)

diff --git a/internal/jptrace/aggregator.go b/internal/jptrace/aggregator.go
index c0af2ebd4bd..1432fb7a5b6 100644
--- a/internal/jptrace/aggregator.go
+++ b/internal/jptrace/aggregator.go
@@ -24,7 +24,6 @@ func AggregateTracesWithLimit(tracesSeq iter.Seq2[[]ptrace.Traces, error], maxSi
 		currentTrace := ptrace.NewTraces()
 		currentTraceID := pcommon.NewTraceIDEmpty()
 		spanCount := 0
-		skipCurrentTrace := false
 
 		tracesSeq(func(traces []ptrace.Traces, err error) bool {
 			if err != nil {
@@ -35,12 +34,9 @@ func AggregateTracesWithLimit(tracesSeq iter.Seq2[[]ptrace.Traces, error], maxSi
 				resources := trace.ResourceSpans()
 				traceID := resources.At(0).ScopeSpans().At(0).Spans().At(0).TraceID()
 				if currentTraceID == traceID {
-					if !skipCurrentTrace {
-						spanCount = mergeTraces(trace, currentTrace, maxSize, spanCount)
-						if truncated {
-							markTraceTruncated(currentTrace)
-							skipCurrentTrace = true
-						}
+					truncated := mergeTraces(trace, currentTrace, maxSize, &spanCount)
+					if truncated {
+						markTraceTruncated(currentTrace)
 					}
 				} else {
 					if currentTrace.ResourceSpans().Len() > 0 {
@@ -50,14 +46,12 @@ func AggregateTracesWithLimit(tracesSeq iter.Seq2[[]ptrace.Traces, error], maxSi
 					}
 					currentTrace = trace
 					currentTraceID = traceID
-					spanCount = countSpans(trace)
-					skipCurrentTrace = false
+					spanCount = trace.SpanCount()
 					if maxSize > 0 && spanCount > maxSize {
 						currentTrace = ptrace.NewTraces()
 						copySpansUpToLimit(trace, currentTrace, maxSize)
 						spanCount = maxSize
 						markTraceTruncated(currentTrace)
-						skipCurrentTrace = true
 					}
 				}
 			}
@@ -70,19 +64,14 @@ func AggregateTracesWithLimit(tracesSeq iter.Seq2[[]ptrace.Traces, error], maxSi
 }
 
 func mergeTraces(src, dest ptrace.Traces, maxSize int, spanCount *int) bool {
-	if maxSize <= 0 {
-		// No limit, merge all
-		resources := src.ResourceSpans()
-		for i := 0; i < resources.Len(); i++ {
-			resource := resources.At(i)
-			resource.CopyTo(dest.ResourceSpans().AppendEmpty())
-		}
-		return false
+	// early exit if already at max
+	if maxSize > 0 && *spanCount >= maxSize {
+		return true
 	}
 
-	// with limit
-	incomingCount := countSpans(src)
-	if *spanCount+incomingCount <= maxSize {
+	incomingCount := src.SpanCount()
+	// check if we can merge all spans without exceeding limit
+	if maxSize <= 0 || *spanCount+incomingCount <= maxSize {
 		resources := src.ResourceSpans()
 		for i := 0; i < resources.Len(); i++ {
 			resource := resources.At(i)
@@ -101,18 +90,6 @@ func mergeTraces(src, dest ptrace.Traces, maxSize int, spanCount *int) bool {
 	return true
 }
 
-func countSpans(trace ptrace.Traces) int {
-	count := 0
-	resources := trace.ResourceSpans()
-	for i := 0; i < resources.Len(); i++ {
-		scopes := resources.At(i).ScopeSpans()
-		for j := 0; j < scopes.Len(); j++ {
-			count += scopes.At(j).Spans().Len()
-		}
-	}
-	return count
-}
-
 func copySpansUpToLimit(src, dest ptrace.Traces, limit int) {
 	copied := 0
 	resources := src.ResourceSpans()
@@ -140,37 +117,22 @@ func copySpansUpToLimit(src, dest ptrace.Traces, limit int) {
 }
 
 func markTraceTruncated(trace ptrace.Traces) {
-	resources := trace.ResourceSpans()
-	if resources.Len() == 0 {
-		return
-	}
-	scopes := resources.At(0).ScopeSpans()
-	if scopes.Len() == 0 {
-		return
-	}
-	spans := scopes.At(0).Spans()
-	if spans.Len() == 0 {
-		return
-	}
-	firstSpan := spans.At(0)
+	// direct access to first span (if truncated, it must exist)
+	firstSpan := trace.ResourceSpans().At(0).ScopeSpans().At(0).Spans().At(0)
 	firstSpan.Attributes().PutBool(truncationMarkerKey, true)
 }
 
 // check if a trace has marked as truncated
 func IsTraceTruncated(trace ptrace.Traces) bool {
-	resources := trace.ResourceSpans()
-	if resources.Len() == 0 {
-		return false
-	}
-	scopes := resources.At(0).ScopeSpans()
-	if scopes.Len() == 0 {
-		return false
-	}
-	spans := scopes.At(0).Spans()
-	if spans.Len() == 0 {
-		return false
+	for i := 0; i < trace.ResourceSpans().Len(); i++ {
+		scopes := trace.ResourceSpans().At(i).ScopeSpans()
+		for j := 0; j < scopes.Len(); j++ {
+			spans := scopes.At(j).Spans()
+			if spans.Len() > 0 {
+				val, exists := spans.At(0).Attributes().Get(truncationMarkerKey)
+				return exists && val.Bool()
+			}
+		}
 	}
-	firstSpan := spans.At(0)
-	val, exists := firstSpan.Attributes().Get(truncationMarkerKey)
-	return exists && val.Bool()
+	return false
 }
diff --git a/internal/jptrace/aggregator_test.go b/internal/jptrace/aggregator_test.go
index 620794b9433..f6fbc2f06dd 100644
--- a/internal/jptrace/aggregator_test.go
+++ b/internal/jptrace/aggregator_test.go
@@ -172,23 +172,12 @@ func TestAggregateTracesWithLimit(t *testing.T) {
 			})
 
 			require.Len(t, result, 1)
-			assert.Equal(t, tt.expectedSpans, countSpans(result[0]))
+			assert.Equal(t, tt.expectedSpans, result[0].SpanCount())
 			assert.Equal(t, tt.expectTruncate, IsTraceTruncated(result[0]))
 		})
 	}
 }
 
-func TestCountSpans(t *testing.T) {
-	trace := ptrace.NewTraces()
-	r1 := trace.ResourceSpans().AppendEmpty()
-	r1.ScopeSpans().AppendEmpty().Spans().AppendEmpty()
-	r1.ScopeSpans().AppendEmpty().Spans().AppendEmpty()
-	r2 := trace.ResourceSpans().AppendEmpty()
-	r2.ScopeSpans().AppendEmpty().Spans().AppendEmpty()
-
-	assert.Equal(t, 3, countSpans(trace))
-}
-
 func TestCopySpansUpToLimit(t *testing.T) {
 	src := ptrace.NewTraces()
 	spans := src.ResourceSpans().AppendEmpty().ScopeSpans().AppendEmpty().Spans()
@@ -199,7 +188,7 @@ func TestCopySpansUpToLimit(t *testing.T) {
 	dest := ptrace.NewTraces()
 	copySpansUpToLimit(src, dest, 3)
 
-	assert.Equal(t, 3, countSpans(dest))
+	assert.Equal(t, 3, dest.SpanCount())
 }
 
 func TestMarkAndCheckTruncated(t *testing.T) {

From 840ebc9788bfc60d0831194d0d98d8ba8dd89531 Mon Sep 17 00:00:00 2001
From: Parship Chowdhury <parshipchowdhury@gmail.com>
Date: Thu, 1 Jan 2026 05:42:47 +0000
Subject: [PATCH 04/22] 1. dry the service tests 2. use range with .All() in
 copySpansUpToLimit 3. argument order to (dest, src, ...)

Signed-off-by: Parship Chowdhury <parshipchowdhury@gmail.com>
---
 .../querysvc/v2/querysvc/service_test.go      | 260 ++++++++----------
 internal/jptrace/aggregator.go                |  22 +-
 internal/jptrace/aggregator_test.go           |   2 +-
 3 files changed, 123 insertions(+), 161 deletions(-)

diff --git a/cmd/jaeger/internal/extension/jaegerquery/internal/querysvc/v2/querysvc/service_test.go b/cmd/jaeger/internal/extension/jaegerquery/internal/querysvc/v2/querysvc/service_test.go
index 34779b3f977..c67dff0eb9c 100644
--- a/cmd/jaeger/internal/extension/jaegerquery/internal/querysvc/v2/querysvc/service_test.go
+++ b/cmd/jaeger/internal/extension/jaegerquery/internal/querysvc/v2/querysvc/service_test.go
@@ -624,161 +624,125 @@ func TestGetCapabilities(t *testing.T) {
 	}
 }
 
-func TestMaxTraceSize_UnderLimit(t *testing.T) {
-	// 3 spans
-	trace := ptrace.NewTraces()
-	resources := trace.ResourceSpans().AppendEmpty()
-	scopes := resources.ScopeSpans().AppendEmpty()
-	for i := 0; i < 3; i++ {
-		span := scopes.Spans().AppendEmpty()
-		span.SetTraceID(testTraceID)
-		span.SetSpanID(pcommon.SpanID([8]byte{byte(i + 1)}))
-		span.SetName(fmt.Sprintf("span-%d", i))
-	}
-
-	responseIter := iter.Seq2[[]ptrace.Traces, error](func(yield func([]ptrace.Traces, error) bool) {
-		yield([]ptrace.Traces{trace}, nil)
-	})
-
-	traceReader := &tracestoremocks.Reader{}
-	dependencyStorage := &depstoremocks.Reader{}
-	options := QueryServiceOptions{
-		MaxTraceSize: 5, // limit is 5, but trace has only 3 spans
-	}
-	tqs := &testQueryService{}
-	tqs.queryService = NewQueryService(traceReader, dependencyStorage, options)
-
-	params := GetTraceParams{
-		TraceIDs: []tracestore.GetTraceParams{{TraceID: testTraceID}},
-	}
-	traceReader.On("GetTraces", mock.Anything, params.TraceIDs).
-		Return(responseIter).Once()
-
-	getTracesIter := tqs.queryService.GetTraces(context.Background(), params)
-	gotTraces, err := jiter.FlattenWithErrors(getTracesIter)
-	require.NoError(t, err)
-	require.Len(t, gotTraces, 1)
-
-	gotSpans := gotTraces[0].ResourceSpans().At(0).ScopeSpans().At(0).Spans()
-	require.Equal(t, 3, gotSpans.Len())
-
-	// no warning should be present
-	firstSpan := gotSpans.At(0)
-	_, hasWarning := firstSpan.Attributes().Get("@jaeger@warnings")
-	require.False(t, hasWarning)
-}
-
-func TestMaxTraceSize_OverLimit(t *testing.T) {
-	// 5 spans split across 2 batches
-	trace1 := ptrace.NewTraces()
-	resources1 := trace1.ResourceSpans().AppendEmpty()
-	scopes1 := resources1.ScopeSpans().AppendEmpty()
-	for i := 0; i < 3; i++ {
-		span := scopes1.Spans().AppendEmpty()
-		span.SetTraceID(testTraceID)
-		span.SetSpanID(pcommon.SpanID([8]byte{byte(i + 1)}))
-		span.SetName(fmt.Sprintf("span-%d", i))
-	}
-
-	trace2 := ptrace.NewTraces()
-	resources2 := trace2.ResourceSpans().AppendEmpty()
-	scopes2 := resources2.ScopeSpans().AppendEmpty()
-	for i := 3; i < 5; i++ {
-		span := scopes2.Spans().AppendEmpty()
-		span.SetTraceID(testTraceID)
-		span.SetSpanID(pcommon.SpanID([8]byte{byte(i + 1)}))
-		span.SetName(fmt.Sprintf("span-%d", i))
-	}
-
-	responseIter := iter.Seq2[[]ptrace.Traces, error](func(yield func([]ptrace.Traces, error) bool) {
-		if !yield([]ptrace.Traces{trace1}, nil) {
-			return
-		}
-		yield([]ptrace.Traces{trace2}, nil)
-	})
-
-	traceReader := &tracestoremocks.Reader{}
-	dependencyStorage := &depstoremocks.Reader{}
-	options := QueryServiceOptions{
-		MaxTraceSize: 3, // Limit is 3, but trace has 5 spans total
-	}
-	tqs := &testQueryService{}
-	tqs.queryService = NewQueryService(traceReader, dependencyStorage, options)
-
-	params := GetTraceParams{
-		TraceIDs: []tracestore.GetTraceParams{{TraceID: testTraceID}},
+func TestMaxTraceSize(t *testing.T) {
+	tests := []struct {
+		name           string
+		maxTraceSize   int
+		createTraces   func() []ptrace.Traces
+		expectedSpans  int
+		expectWarning  bool
+		warningPattern string
+	}{
+		{
+			name:         "under_limit",
+			maxTraceSize: 5,
+			createTraces: func() []ptrace.Traces {
+				trace := ptrace.NewTraces()
+				scopes := trace.ResourceSpans().AppendEmpty().ScopeSpans().AppendEmpty()
+				for i := 0; i < 3; i++ {
+					span := scopes.Spans().AppendEmpty()
+					span.SetTraceID(testTraceID)
+					span.SetSpanID(pcommon.SpanID([8]byte{byte(i + 1)}))
+					span.SetName(fmt.Sprintf("span-%d", i))
+				}
+				return []ptrace.Traces{trace}
+			},
+			expectedSpans: 3,
+			expectWarning: false,
+		},
+		{
+			name:         "over_limit",
+			maxTraceSize: 3,
+			createTraces: func() []ptrace.Traces {
+				trace1 := ptrace.NewTraces()
+				scopes1 := trace1.ResourceSpans().AppendEmpty().ScopeSpans().AppendEmpty()
+				for i := 0; i < 3; i++ {
+					span := scopes1.Spans().AppendEmpty()
+					span.SetTraceID(testTraceID)
+					span.SetSpanID(pcommon.SpanID([8]byte{byte(i + 1)}))
+					span.SetName(fmt.Sprintf("span-%d", i))
+				}
+
+				trace2 := ptrace.NewTraces()
+				scopes2 := trace2.ResourceSpans().AppendEmpty().ScopeSpans().AppendEmpty()
+				for i := 3; i < 5; i++ {
+					span := scopes2.Spans().AppendEmpty()
+					span.SetTraceID(testTraceID)
+					span.SetSpanID(pcommon.SpanID([8]byte{byte(i + 1)}))
+					span.SetName(fmt.Sprintf("span-%d", i))
+				}
+				return []ptrace.Traces{trace1, trace2}
+			},
+			expectedSpans:  3,
+			expectWarning:  true,
+			warningPattern: "trace has more than 3 spans",
+		},
+		{
+			name:         "exactly_at_limit",
+			maxTraceSize: 3,
+			createTraces: func() []ptrace.Traces {
+				trace := ptrace.NewTraces()
+				scopes := trace.ResourceSpans().AppendEmpty().ScopeSpans().AppendEmpty()
+				for i := 0; i < 3; i++ {
+					span := scopes.Spans().AppendEmpty()
+					span.SetTraceID(testTraceID)
+					span.SetSpanID(pcommon.SpanID([8]byte{byte(i + 1)}))
+					span.SetName(fmt.Sprintf("span-%d", i))
+				}
+				return []ptrace.Traces{trace}
+			},
+			expectedSpans: 3,
+			expectWarning: false,
+		},
 	}
-	traceReader.On("GetTraces", mock.Anything, params.TraceIDs).
-		Return(responseIter).Once()
 
-	getTracesIter := tqs.queryService.GetTraces(context.Background(), params)
-	gotTraces, err := jiter.FlattenWithErrors(getTracesIter)
-	require.NoError(t, err)
-	require.Len(t, gotTraces, 1)
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			traces := tt.createTraces()
+			responseIter := iter.Seq2[[]ptrace.Traces, error](func(yield func([]ptrace.Traces, error) bool) {
+				for _, trace := range traces {
+					if !yield([]ptrace.Traces{trace}, nil) {
+						return
+					}
+				}
+			})
 
-	// count total spans in the result
-	totalSpans := 0
-	resources := gotTraces[0].ResourceSpans()
-	for i := 0; i < resources.Len(); i++ {
-		scopes := resources.At(i).ScopeSpans()
-		for j := 0; j < scopes.Len(); j++ {
-			totalSpans += scopes.At(j).Spans().Len()
-		}
-	}
+			traceReader := &tracestoremocks.Reader{}
+			dependencyStorage := &depstoremocks.Reader{}
+			options := QueryServiceOptions{
+				MaxTraceSize: tt.maxTraceSize,
+			}
+			tqs := &testQueryService{}
+			tqs.queryService = NewQueryService(traceReader, dependencyStorage, options)
 
-	// Only 3 spans should be present(the limit)
-	require.Equal(t, 3, totalSpans)
-
-	// there should be a warning for the first span
-	firstSpan := resources.At(0).ScopeSpans().At(0).Spans().At(0)
-	warningsAttr, hasWarning := firstSpan.Attributes().Get("@jaeger@warnings")
-	require.True(t, hasWarning)
-	require.Equal(t, pcommon.ValueTypeSlice, warningsAttr.Type())
-	warnings := warningsAttr.Slice()
-	require.Positive(t, warnings.Len())
-	require.Contains(t, warnings.At(warnings.Len()-1).Str(), "trace has more than 3 spans")
-}
+			params := GetTraceParams{
+				TraceIDs: []tracestore.GetTraceParams{{TraceID: testTraceID}},
+			}
+			traceReader.On("GetTraces", mock.Anything, params.TraceIDs).
+				Return(responseIter).Once()
 
-func TestMaxTraceSize_ExactlyAtLimit(t *testing.T) {
-	// 3 spans
-	trace := ptrace.NewTraces()
-	resources := trace.ResourceSpans().AppendEmpty()
-	scopes := resources.ScopeSpans().AppendEmpty()
-	for i := 0; i < 3; i++ {
-		span := scopes.Spans().AppendEmpty()
-		span.SetTraceID(testTraceID)
-		span.SetSpanID(pcommon.SpanID([8]byte{byte(i + 1)}))
-		span.SetName(fmt.Sprintf("span-%d", i))
-	}
+			getTracesIter := tqs.queryService.GetTraces(context.Background(), params)
+			gotTraces, err := jiter.FlattenWithErrors(getTracesIter)
+			require.NoError(t, err)
+			require.Len(t, gotTraces, 1)
 
-	responseIter := iter.Seq2[[]ptrace.Traces, error](func(yield func([]ptrace.Traces, error) bool) {
-		yield([]ptrace.Traces{trace}, nil)
-	})
+			// count total spans
+			actualSpans := gotTraces[0].SpanCount()
+			require.Equal(t, tt.expectedSpans, actualSpans)
 
-	traceReader := &tracestoremocks.Reader{}
-	dependencyStorage := &depstoremocks.Reader{}
-	options := QueryServiceOptions{
-		MaxTraceSize: 3, // Limit is exactly 3, trace has 3 spans
-	}
-	tqs := &testQueryService{}
-	tqs.queryService = NewQueryService(traceReader, dependencyStorage, options)
+			// check warning
+			firstSpan := gotTraces[0].ResourceSpans().At(0).ScopeSpans().At(0).Spans().At(0)
+			warningsAttr, hasWarning := firstSpan.Attributes().Get("@jaeger@warnings")
 
-	params := GetTraceParams{
-		TraceIDs: []tracestore.GetTraceParams{{TraceID: testTraceID}},
+			if tt.expectWarning {
+				require.True(t, hasWarning, "expected warning but none found")
+				require.Equal(t, pcommon.ValueTypeSlice, warningsAttr.Type())
+				warnings := warningsAttr.Slice()
+				require.Positive(t, warnings.Len())
+				require.Contains(t, warnings.At(warnings.Len()-1).Str(), tt.warningPattern)
+			} else {
+				require.False(t, hasWarning, "unexpected warning found")
+			}
+		})
 	}
-	traceReader.On("GetTraces", mock.Anything, params.TraceIDs).
-		Return(responseIter).Once()
-
-	getTracesIter := tqs.queryService.GetTraces(context.Background(), params)
-	gotTraces, err := jiter.FlattenWithErrors(getTracesIter)
-	require.NoError(t, err)
-	require.Len(t, gotTraces, 1)
-
-	gotSpans := gotTraces[0].ResourceSpans().At(0).ScopeSpans().At(0).Spans()
-	// all 3 spans should be present
-	require.Equal(t, 3, gotSpans.Len())
-
-	firstSpan := gotSpans.At(0)
-	_, hasWarning := firstSpan.Attributes().Get("@jaeger@warnings")
-	require.False(t, hasWarning)
 }
diff --git a/internal/jptrace/aggregator.go b/internal/jptrace/aggregator.go
index 1432fb7a5b6..2627e3d73cf 100644
--- a/internal/jptrace/aggregator.go
+++ b/internal/jptrace/aggregator.go
@@ -49,7 +49,7 @@ func AggregateTracesWithLimit(tracesSeq iter.Seq2[[]ptrace.Traces, error], maxSi
 					spanCount = trace.SpanCount()
 					if maxSize > 0 && spanCount > maxSize {
 						currentTrace = ptrace.NewTraces()
-						copySpansUpToLimit(trace, currentTrace, maxSize)
+						copySpansUpToLimit(currentTrace, trace, maxSize)
 						spanCount = maxSize
 						markTraceTruncated(currentTrace)
 					}
@@ -84,32 +84,30 @@ func mergeTraces(src, dest ptrace.Traces, maxSize int, spanCount *int) bool {
 	// partial copy
 	remaining := maxSize - *spanCount
 	if remaining > 0 {
-		copySpansUpToLimit(src, dest, remaining)
+		copySpansUpToLimit(dest, src, remaining)
 		*spanCount = maxSize
 	}
 	return true
 }
 
-func copySpansUpToLimit(src, dest ptrace.Traces, limit int) {
+func copySpansUpToLimit(dest, src ptrace.Traces, limit int) {
 	copied := 0
-	resources := src.ResourceSpans()
 
-	for i := 0; i < resources.Len() && copied < limit; i++ {
-		srcResource := resources.At(i)
+	for _, srcResource := range src.ResourceSpans().All() {
 		destResource := dest.ResourceSpans().AppendEmpty()
 		srcResource.Resource().CopyTo(destResource.Resource())
 		destResource.SetSchemaUrl(srcResource.SchemaUrl())
 
-		scopes := srcResource.ScopeSpans()
-		for j := 0; j < scopes.Len() && copied < limit; j++ {
-			srcScope := scopes.At(j)
+		for _, srcScope := range srcResource.ScopeSpans().All() {
 			destScope := destResource.ScopeSpans().AppendEmpty()
 			srcScope.Scope().CopyTo(destScope.Scope())
 			destScope.SetSchemaUrl(srcScope.SchemaUrl())
 
-			spans := srcScope.Spans()
-			for k := 0; k < spans.Len() && copied < limit; k++ {
-				spans.At(k).CopyTo(destScope.Spans().AppendEmpty())
+			for _, span := range srcScope.Spans().All() {
+				if copied >= limit {
+					return
+				}
+				span.CopyTo(destScope.Spans().AppendEmpty())
 				copied++
 			}
 		}
diff --git a/internal/jptrace/aggregator_test.go b/internal/jptrace/aggregator_test.go
index f6fbc2f06dd..6c12794c442 100644
--- a/internal/jptrace/aggregator_test.go
+++ b/internal/jptrace/aggregator_test.go
@@ -186,7 +186,7 @@ func TestCopySpansUpToLimit(t *testing.T) {
 	}
 
 	dest := ptrace.NewTraces()
-	copySpansUpToLimit(src, dest, 3)
+	copySpansUpToLimit(dest, src, 3)
 
 	assert.Equal(t, 3, dest.SpanCount())
 }

From 5b73044d331a07b528f5fdc89b306e1f39fd28f0 Mon Sep 17 00:00:00 2001
From: Parship Chowdhury <parshipchowdhury@gmail.com>
Date: Thu, 1 Jan 2026 05:57:59 +0000
Subject: [PATCH 05/22] added a comment for TestMaxTraceSize

Signed-off-by: Parship Chowdhury <parshipchowdhury@gmail.com>
---
 .../jaegerquery/internal/querysvc/v2/querysvc/service_test.go    | 1 +
 1 file changed, 1 insertion(+)

diff --git a/cmd/jaeger/internal/extension/jaegerquery/internal/querysvc/v2/querysvc/service_test.go b/cmd/jaeger/internal/extension/jaegerquery/internal/querysvc/v2/querysvc/service_test.go
index c67dff0eb9c..76f33abf4e0 100644
--- a/cmd/jaeger/internal/extension/jaegerquery/internal/querysvc/v2/querysvc/service_test.go
+++ b/cmd/jaeger/internal/extension/jaegerquery/internal/querysvc/v2/querysvc/service_test.go
@@ -624,6 +624,7 @@ func TestGetCapabilities(t *testing.T) {
 	}
 }
 
+// Consolidate Underlimit, Overlimit and Exactly at limit tests
 func TestMaxTraceSize(t *testing.T) {
 	tests := []struct {
 		name           string

From 9fa50a27fb57245be3676d3491cc4ae463d02bf3 Mon Sep 17 00:00:00 2001
From: Parship Chowdhury <parshipchowdhury@gmail.com>
Date: Fri, 2 Jan 2026 10:14:14 +0530
Subject: [PATCH 06/22] Update internal/jptrace/aggregator.go

Co-authored-by: Yuri Shkuro <yurishkuro@users.noreply.github.com>
Signed-off-by: Parship Chowdhury <parshipchowdhury@gmail.com>
---
 internal/jptrace/aggregator.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/internal/jptrace/aggregator.go b/internal/jptrace/aggregator.go
index 2627e3d73cf..43c34ece62f 100644
--- a/internal/jptrace/aggregator.go
+++ b/internal/jptrace/aggregator.go
@@ -63,7 +63,7 @@ func AggregateTracesWithLimit(tracesSeq iter.Seq2[[]ptrace.Traces, error], maxSi
 	}
 }
 
-func mergeTraces(src, dest ptrace.Traces, maxSize int, spanCount *int) bool {
+func mergeTraces(dest, src ptrace.Traces, maxSize int, spanCount *int) bool {
 	// early exit if already at max
 	if maxSize > 0 && *spanCount >= maxSize {
 		return true

From 5a7614c6454ce9439c2ed9207e1bb000b2fb7a3f Mon Sep 17 00:00:00 2001
From: Parship Chowdhury <parshipchowdhury@gmail.com>
Date: Fri, 2 Jan 2026 10:14:52 +0530
Subject: [PATCH 07/22] Update internal/jptrace/aggregator.go

Co-authored-by: Yuri Shkuro <yurishkuro@users.noreply.github.com>
Signed-off-by: Parship Chowdhury <parshipchowdhury@gmail.com>
---
 internal/jptrace/aggregator.go | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/internal/jptrace/aggregator.go b/internal/jptrace/aggregator.go
index 43c34ece62f..bb60621ea7d 100644
--- a/internal/jptrace/aggregator.go
+++ b/internal/jptrace/aggregator.go
@@ -125,8 +125,7 @@ func IsTraceTruncated(trace ptrace.Traces) bool {
 	for i := 0; i < trace.ResourceSpans().Len(); i++ {
 		scopes := trace.ResourceSpans().At(i).ScopeSpans()
 		for j := 0; j < scopes.Len(); j++ {
-			spans := scopes.At(j).Spans()
-			if spans.Len() > 0 {
+			for _, span := range scope.Spans().All()
 				val, exists := spans.At(0).Attributes().Get(truncationMarkerKey)
 				return exists && val.Bool()
 			}

From 928ae3f9f692d17cb1ebe915428689f05d19100a Mon Sep 17 00:00:00 2001
From: Parship Chowdhury <parshipchowdhury@gmail.com>
Date: Fri, 2 Jan 2026 05:17:48 +0000
Subject: [PATCH 08/22] move markTraceTruncated inside mergeTraces, use
 mergeTraces instead of copySpansUpToLimit, fix early return

Signed-off-by: Parship Chowdhury <parshipchowdhury@gmail.com>
---
 internal/jptrace/aggregator.go | 26 +++++++++++++++-----------
 1 file changed, 15 insertions(+), 11 deletions(-)

diff --git a/internal/jptrace/aggregator.go b/internal/jptrace/aggregator.go
index bb60621ea7d..75df379117b 100644
--- a/internal/jptrace/aggregator.go
+++ b/internal/jptrace/aggregator.go
@@ -34,10 +34,7 @@ func AggregateTracesWithLimit(tracesSeq iter.Seq2[[]ptrace.Traces, error], maxSi
 				resources := trace.ResourceSpans()
 				traceID := resources.At(0).ScopeSpans().At(0).Spans().At(0).TraceID()
 				if currentTraceID == traceID {
-					truncated := mergeTraces(trace, currentTrace, maxSize, &spanCount)
-					if truncated {
-						markTraceTruncated(currentTrace)
-					}
+					mergeTraces(currentTrace, trace, maxSize, &spanCount)
 				} else {
 					if currentTrace.ResourceSpans().Len() > 0 {
 						if !yield(currentTrace, nil) {
@@ -49,9 +46,8 @@ func AggregateTracesWithLimit(tracesSeq iter.Seq2[[]ptrace.Traces, error], maxSi
 					spanCount = trace.SpanCount()
 					if maxSize > 0 && spanCount > maxSize {
 						currentTrace = ptrace.NewTraces()
-						copySpansUpToLimit(currentTrace, trace, maxSize)
-						spanCount = maxSize
-						markTraceTruncated(currentTrace)
+						spanCount = 0
+						mergeTraces(currentTrace, trace, maxSize, &spanCount)
 					}
 				}
 			}
@@ -66,6 +62,7 @@ func AggregateTracesWithLimit(tracesSeq iter.Seq2[[]ptrace.Traces, error], maxSi
 func mergeTraces(dest, src ptrace.Traces, maxSize int, spanCount *int) bool {
 	// early exit if already at max
 	if maxSize > 0 && *spanCount >= maxSize {
+		markTraceTruncated(dest)
 		return true
 	}
 
@@ -87,6 +84,7 @@ func mergeTraces(dest, src ptrace.Traces, maxSize int, spanCount *int) bool {
 		copySpansUpToLimit(dest, src, remaining)
 		*spanCount = maxSize
 	}
+	markTraceTruncated(dest)
 	return true
 }
 
@@ -94,11 +92,17 @@ func copySpansUpToLimit(dest, src ptrace.Traces, limit int) {
 	copied := 0
 
 	for _, srcResource := range src.ResourceSpans().All() {
+		if copied >= limit {
+			return
+		}
 		destResource := dest.ResourceSpans().AppendEmpty()
 		srcResource.Resource().CopyTo(destResource.Resource())
 		destResource.SetSchemaUrl(srcResource.SchemaUrl())
 
 		for _, srcScope := range srcResource.ScopeSpans().All() {
+			if copied >= limit {
+				return
+			}
 			destScope := destResource.ScopeSpans().AppendEmpty()
 			srcScope.Scope().CopyTo(destScope.Scope())
 			destScope.SetSchemaUrl(srcScope.SchemaUrl())
@@ -122,10 +126,10 @@ func markTraceTruncated(trace ptrace.Traces) {
 
 // check if a trace has marked as truncated
 func IsTraceTruncated(trace ptrace.Traces) bool {
-	for i := 0; i < trace.ResourceSpans().Len(); i++ {
-		scopes := trace.ResourceSpans().At(i).ScopeSpans()
-		for j := 0; j < scopes.Len(); j++ {
-			for _, span := range scope.Spans().All()
+	for _, resource := range trace.ResourceSpans().All() {
+		for _, scope := range resource.ScopeSpans().All() {
+			spans := scope.Spans()
+			if spans.Len() > 0 {
 				val, exists := spans.At(0).Attributes().Get(truncationMarkerKey)
 				return exists && val.Bool()
 			}

From 78f1aaa0de6fd7ddcf071e2ffb108dde55809cf2 Mon Sep 17 00:00:00 2001
From: Parship Chowdhury <parshipchowdhury@gmail.com>
Date: Thu, 8 Jan 2026 08:52:18 +0000
Subject: [PATCH 09/22] simplify

Signed-off-by: Parship Chowdhury <parshipchowdhury@gmail.com>
---
 .../internal/querysvc/v2/querysvc/service.go  | 29 -------------------
 internal/jptrace/aggregator.go                | 26 ++++-------------
 internal/jptrace/aggregator_test.go           | 22 ++++++++++----
 3 files changed, 22 insertions(+), 55 deletions(-)

diff --git a/cmd/jaeger/internal/extension/jaegerquery/internal/querysvc/v2/querysvc/service.go b/cmd/jaeger/internal/extension/jaegerquery/internal/querysvc/v2/querysvc/service.go
index 313987a75ca..ad94d7e16e7 100644
--- a/cmd/jaeger/internal/extension/jaegerquery/internal/querysvc/v2/querysvc/service.go
+++ b/cmd/jaeger/internal/extension/jaegerquery/internal/querysvc/v2/querysvc/service.go
@@ -6,7 +6,6 @@ package querysvc
 import (
 	"context"
 	"errors"
-	"fmt"
 	"iter"
 	"time"
 
@@ -205,37 +204,9 @@ func (qs QueryService) receiveTraces(
 		seq(processTraces)
 	} else {
 		jptrace.AggregateTracesWithLimit(seq, qs.options.MaxTraceSize)(func(trace ptrace.Traces, err error) bool {
-			// Add warning if trace was truncated
-			if err == nil && qs.options.MaxTraceSize > 0 && jptrace.IsTraceTruncated(trace) {
-				qs.addTruncationWarning(trace)
-			}
 			return processTraces([]ptrace.Traces{trace}, err)
 		})
 	}
 
 	return foundTraceIDs, proceed
 }
-
-// add a warning to the first span of the trace
-func (qs QueryService) addTruncationWarning(trace ptrace.Traces) {
-	resources := trace.ResourceSpans()
-	if resources.Len() == 0 {
-		return
-	}
-
-	scopes := resources.At(0).ScopeSpans()
-	if scopes.Len() == 0 {
-		return
-	}
-
-	spans := scopes.At(0).Spans()
-	if spans.Len() == 0 {
-		return
-	}
-
-	firstSpan := spans.At(0)
-	firstSpan.Attributes().Remove("@jaeger@truncated")
-	jptrace.AddWarnings(firstSpan,
-		fmt.Sprintf("trace has more than %d spans, showing first %d spans only",
-			qs.options.MaxTraceSize, qs.options.MaxTraceSize))
-}
diff --git a/internal/jptrace/aggregator.go b/internal/jptrace/aggregator.go
index 75df379117b..c3c4a0b4196 100644
--- a/internal/jptrace/aggregator.go
+++ b/internal/jptrace/aggregator.go
@@ -4,14 +4,13 @@
 package jptrace
 
 import (
+	"fmt"
 	"iter"
 
 	"go.opentelemetry.io/collector/pdata/pcommon"
 	"go.opentelemetry.io/collector/pdata/ptrace"
 )
 
-const truncationMarkerKey = "@jaeger@truncated"
-
 // AggregateTraces aggregates a sequence of trace batches into individual traces.
 //
 // The `tracesSeq` input must adhere to the chunking requirements of tracestore.Reader.GetTraces.
@@ -62,7 +61,7 @@ func AggregateTracesWithLimit(tracesSeq iter.Seq2[[]ptrace.Traces, error], maxSi
 func mergeTraces(dest, src ptrace.Traces, maxSize int, spanCount *int) bool {
 	// early exit if already at max
 	if maxSize > 0 && *spanCount >= maxSize {
-		markTraceTruncated(dest)
+		markTraceTruncated(dest, maxSize)
 		return true
 	}
 
@@ -84,7 +83,7 @@ func mergeTraces(dest, src ptrace.Traces, maxSize int, spanCount *int) bool {
 		copySpansUpToLimit(dest, src, remaining)
 		*spanCount = maxSize
 	}
-	markTraceTruncated(dest)
+	markTraceTruncated(dest, maxSize)
 	return true
 }
 
@@ -118,22 +117,9 @@ func copySpansUpToLimit(dest, src ptrace.Traces, limit int) {
 	}
 }
 
-func markTraceTruncated(trace ptrace.Traces) {
+func markTraceTruncated(trace ptrace.Traces, maxSize int) {
 	// direct access to first span (if truncated, it must exist)
 	firstSpan := trace.ResourceSpans().At(0).ScopeSpans().At(0).Spans().At(0)
-	firstSpan.Attributes().PutBool(truncationMarkerKey, true)
-}
-
-// check if a trace has marked as truncated
-func IsTraceTruncated(trace ptrace.Traces) bool {
-	for _, resource := range trace.ResourceSpans().All() {
-		for _, scope := range resource.ScopeSpans().All() {
-			spans := scope.Spans()
-			if spans.Len() > 0 {
-				val, exists := spans.At(0).Attributes().Get(truncationMarkerKey)
-				return exists && val.Bool()
-			}
-		}
-	}
-	return false
+	AddWarnings(firstSpan,
+		fmt.Sprintf("trace has more than %d spans, showing first %d spans only", maxSize, maxSize))
 }
diff --git a/internal/jptrace/aggregator_test.go b/internal/jptrace/aggregator_test.go
index 6c12794c442..ea872add4c1 100644
--- a/internal/jptrace/aggregator_test.go
+++ b/internal/jptrace/aggregator_test.go
@@ -4,6 +4,7 @@
 package jptrace
 
 import (
+	"fmt"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -173,7 +174,14 @@ func TestAggregateTracesWithLimit(t *testing.T) {
 
 			require.Len(t, result, 1)
 			assert.Equal(t, tt.expectedSpans, result[0].SpanCount())
-			assert.Equal(t, tt.expectTruncate, IsTraceTruncated(result[0]))
+
+			// Check for truncation warning
+			if tt.expectTruncate {
+				firstSpan := result[0].ResourceSpans().At(0).ScopeSpans().At(0).Spans().At(0)
+				warnings := GetWarnings(firstSpan)
+				assert.NotEmpty(t, warnings, "expected truncation warning")
+				assert.Contains(t, warnings[len(warnings)-1], fmt.Sprintf("trace has more than %d spans", tt.maxSize))
+			}
 		})
 	}
 }
@@ -193,9 +201,11 @@ func TestCopySpansUpToLimit(t *testing.T) {
 
 func TestMarkAndCheckTruncated(t *testing.T) {
 	trace := ptrace.NewTraces()
-	trace.ResourceSpans().AppendEmpty().ScopeSpans().AppendEmpty().Spans().AppendEmpty()
-
-	assert.False(t, IsTraceTruncated(trace))
-	markTraceTruncated(trace)
-	assert.True(t, IsTraceTruncated(trace))
+	firstSpan := trace.ResourceSpans().AppendEmpty().ScopeSpans().AppendEmpty().Spans().AppendEmpty()
+	assert.Empty(t, GetWarnings(firstSpan))
+	markTraceTruncated(trace, 10)
+	// Now should have truncation warning
+	warnings := GetWarnings(firstSpan)
+	assert.NotEmpty(t, warnings)
+	assert.Contains(t, warnings[0], "trace has more than 10 spans")
 }

From 9b4b39e6a25164e57820b59ac3989a6020803e84 Mon Sep 17 00:00:00 2001
From: Parship Chowdhury <parshipchowdhury@gmail.com>
Date: Sat, 31 Jan 2026 07:11:21 +0000
Subject: [PATCH 10/22] fix

Signed-off-by: Parship Chowdhury <parshipchowdhury@gmail.com>
---
 internal/jptrace/aggregator.go | 21 +++++++++------------
 1 file changed, 9 insertions(+), 12 deletions(-)

diff --git a/internal/jptrace/aggregator.go b/internal/jptrace/aggregator.go
index 9588bb167e1..edd3c0487d0 100644
--- a/internal/jptrace/aggregator.go
+++ b/internal/jptrace/aggregator.go
@@ -37,7 +37,7 @@ func AggregateTracesWithLimit(tracesSeq iter.Seq2[[]ptrace.Traces, error], maxSi
 				resources := trace.ResourceSpans()
 				traceID := resources.At(0).ScopeSpans().At(0).Spans().At(0).TraceID()
 				if currentTraceID == traceID {
-					mergeTraces(currentTrace, trace, maxSize, &spanCount)
+					mergeTracesWithLimit(currentTrace, trace, maxSize, &spanCount)
 				} else {
 					if currentTrace.SpanCount() > 0 {
 						if !yield(currentTrace, nil) {
@@ -51,7 +51,7 @@ func AggregateTracesWithLimit(tracesSeq iter.Seq2[[]ptrace.Traces, error], maxSi
 					if maxSize > 0 && spanCount > maxSize {
 						currentTrace = ptrace.NewTraces()
 						spanCount = 0
-						mergeTraces(currentTrace, trace, maxSize, &spanCount)
+						mergeTracesWithLimit(currentTrace, trace, maxSize, &spanCount)
 					}
 				}
 			}
@@ -63,7 +63,7 @@ func AggregateTracesWithLimit(tracesSeq iter.Seq2[[]ptrace.Traces, error], maxSi
 	}
 }
 
-func mergeTraces(dest, src ptrace.Traces, maxSize int, spanCount *int) bool {
+func mergeTracesWithLimit(dest, src ptrace.Traces, maxSize int, spanCount *int) bool {
 	// early exit if already at max
 	if maxSize > 0 && *spanCount >= maxSize {
 		markTraceTruncated(dest, maxSize)
@@ -73,11 +73,7 @@ func mergeTraces(dest, src ptrace.Traces, maxSize int, spanCount *int) bool {
 	incomingCount := src.SpanCount()
 	// check if we can merge all spans without exceeding limit
 	if maxSize <= 0 || *spanCount+incomingCount <= maxSize {
-		resources := src.ResourceSpans()
-		for i := 0; i < resources.Len(); i++ {
-			resource := resources.At(i)
-			resource.CopyTo(dest.ResourceSpans().AppendEmpty())
-		}
+		MergeTraces(dest, src)
 		*spanCount += incomingCount
 		return false
 	}
@@ -133,8 +129,9 @@ func MergeTraces(dest, src ptrace.Traces) {
 }
 
 func markTraceTruncated(trace ptrace.Traces, maxSize int) {
-	// direct access to first span (if truncated, it must exist)
-	firstSpan := trace.ResourceSpans().At(0).ScopeSpans().At(0).Spans().At(0)
-	AddWarnings(firstSpan,
-		fmt.Sprintf("trace has more than %d spans, showing first %d spans only", maxSize, maxSize))
+	SpanIter(trace)(func(_ SpanIterPos, span ptrace.Span) bool {
+		AddWarnings(span,
+			fmt.Sprintf("trace has more than %d spans, showing first %d spans only", maxSize, maxSize))
+		return false // stop after first span
+	})
 }

From b85ca1e71870e72e0e98208cd0bf8453dc14e9cc Mon Sep 17 00:00:00 2001
From: Parship Chowdhury <parshipchowdhury@gmail.com>
Date: Wed, 11 Feb 2026 22:32:33 +0530
Subject: [PATCH 11/22] Update internal/jptrace/aggregator.go

Co-authored-by: Yuri Shkuro <yurishkuro@users.noreply.github.com>
Signed-off-by: Parship Chowdhury <parshipchowdhury@gmail.com>
---
 internal/jptrace/aggregator.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/internal/jptrace/aggregator.go b/internal/jptrace/aggregator.go
index edd3c0487d0..c093ecb5409 100644
--- a/internal/jptrace/aggregator.go
+++ b/internal/jptrace/aggregator.go
@@ -63,7 +63,7 @@ func AggregateTracesWithLimit(tracesSeq iter.Seq2[[]ptrace.Traces, error], maxSi
 	}
 }
 
-func mergeTracesWithLimit(dest, src ptrace.Traces, maxSize int, spanCount *int) bool {
+func mergeTracesWithLimit(dest, src ptrace.Traces, maxSize int, spanCount int) int {
 	// early exit if already at max
 	if maxSize > 0 && *spanCount >= maxSize {
 		markTraceTruncated(dest, maxSize)

From cf22d941d12b00f786a7d5b9f0a5500ad91e23b8 Mon Sep 17 00:00:00 2001
From: Parship Chowdhury <parshipchowdhury@gmail.com>
Date: Thu, 1 Jan 2026 10:10:43 +0530
Subject: [PATCH 12/22] Update internal/jptrace/aggregator.go

Co-authored-by: Yuri Shkuro <yurishkuro@users.noreply.github.com>
Signed-off-by: Parship Chowdhury <parshipchowdhury@gmail.com>
---
 internal/jptrace/aggregator.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/internal/jptrace/aggregator.go b/internal/jptrace/aggregator.go
index c093ecb5409..edd3c0487d0 100644
--- a/internal/jptrace/aggregator.go
+++ b/internal/jptrace/aggregator.go
@@ -63,7 +63,7 @@ func AggregateTracesWithLimit(tracesSeq iter.Seq2[[]ptrace.Traces, error], maxSi
 	}
 }
 
-func mergeTracesWithLimit(dest, src ptrace.Traces, maxSize int, spanCount int) int {
+func mergeTracesWithLimit(dest, src ptrace.Traces, maxSize int, spanCount *int) bool {
 	// early exit if already at max
 	if maxSize > 0 && *spanCount >= maxSize {
 		markTraceTruncated(dest, maxSize)

From de8c332cb1418c1d660472ca187cc17649ee1f6a Mon Sep 17 00:00:00 2001
From: Parship Chowdhury <parshipchowdhury@gmail.com>
Date: Thu, 1 Jan 2026 05:57:59 +0000
Subject: [PATCH 13/22] added a comment for TestMaxTraceSize

Signed-off-by: Parship Chowdhury <parshipchowdhury@gmail.com>
---
 internal/jptrace/aggregator.go | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/internal/jptrace/aggregator.go b/internal/jptrace/aggregator.go
index edd3c0487d0..e0ee5684879 100644
--- a/internal/jptrace/aggregator.go
+++ b/internal/jptrace/aggregator.go
@@ -52,6 +52,9 @@ func AggregateTracesWithLimit(tracesSeq iter.Seq2[[]ptrace.Traces, error], maxSi
 						currentTrace = ptrace.NewTraces()
 						spanCount = 0
 						mergeTracesWithLimit(currentTrace, trace, maxSize, &spanCount)
+						copySpansUpToLimit(currentTrace, trace, maxSize)
+						spanCount = maxSize
+						markTraceTruncated(currentTrace)
 					}
 				}
 			}

From 416be9223eac64b9101f35640371e640d28722f6 Mon Sep 17 00:00:00 2001
From: Parship Chowdhury <parshipchowdhury@gmail.com>
Date: Wed, 11 Feb 2026 19:36:41 +0000
Subject: [PATCH 14/22] fix

Signed-off-by: Parship Chowdhury <parshipchowdhury@gmail.com>
---
 internal/jptrace/aggregator.go | 3 +--
 jaeger-ui                      | 2 +-
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/internal/jptrace/aggregator.go b/internal/jptrace/aggregator.go
index e0ee5684879..1c93bc70d74 100644
--- a/internal/jptrace/aggregator.go
+++ b/internal/jptrace/aggregator.go
@@ -51,10 +51,9 @@ func AggregateTracesWithLimit(tracesSeq iter.Seq2[[]ptrace.Traces, error], maxSi
 					if maxSize > 0 && spanCount > maxSize {
 						currentTrace = ptrace.NewTraces()
 						spanCount = 0
-						mergeTracesWithLimit(currentTrace, trace, maxSize, &spanCount)
 						copySpansUpToLimit(currentTrace, trace, maxSize)
 						spanCount = maxSize
-						markTraceTruncated(currentTrace)
+						markTraceTruncated(currentTrace, maxSize)
 					}
 				}
 			}
diff --git a/jaeger-ui b/jaeger-ui
index ca725430b1d..1ceadb6f6fb 160000
--- a/jaeger-ui
+++ b/jaeger-ui
@@ -1 +1 @@
-Subproject commit ca725430b1d05893ff072abd4389509ec4fe56af
+Subproject commit 1ceadb6f6fb29774bfaa77a6148499ff7a04902b

From 4a9399651b819bcb9ef5d37a250963447b3d903f Mon Sep 17 00:00:00 2001
From: Yuri Shkuro <github@ysh.us>
Date: Thu, 26 Feb 2026 12:13:45 -0500
Subject: [PATCH 15/22] fix ui

Signed-off-by: Yuri Shkuro <github@ysh.us>
---
 jaeger-ui | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/jaeger-ui b/jaeger-ui
index 1ceadb6f6fb..f67531616c7 160000
--- a/jaeger-ui
+++ b/jaeger-ui
@@ -1 +1 @@
-Subproject commit 1ceadb6f6fb29774bfaa77a6148499ff7a04902b
+Subproject commit f67531616c7162ee8b0bb382d69c2d6a006e972c

From d979c25750f8ad9f3787ecef67962f337313e134 Mon Sep 17 00:00:00 2001
From: Yuri Shkuro <yurishkuro@users.noreply.github.com>
Date: Thu, 26 Feb 2026 12:17:18 -0500
Subject: [PATCH 16/22] Apply suggestions from code review

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Signed-off-by: Yuri Shkuro <yurishkuro@users.noreply.github.com>
---
 cmd/jaeger/internal/extension/jaegerquery/internal/flags.go   | 2 +-
 cmd/jaeger/internal/extension/jaegerquery/querysvc/service.go | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/cmd/jaeger/internal/extension/jaegerquery/internal/flags.go b/cmd/jaeger/internal/extension/jaegerquery/internal/flags.go
index aa24c8bfb53..acce6ff1eea 100644
--- a/cmd/jaeger/internal/extension/jaegerquery/internal/flags.go
+++ b/cmd/jaeger/internal/extension/jaegerquery/internal/flags.go
@@ -36,7 +36,7 @@ type QueryOptions struct {
 	Tenancy tenancy.Options `mapstructure:"multi_tenancy"`
 	// MaxClockSkewAdjust is the maximum duration by which jaeger-query will adjust a span.
 	MaxClockSkewAdjust time.Duration `mapstructure:"max_clock_skew_adjust"  valid:"optional"`
-	// MaxTraceSize is the max no. of spans allowed per trace.
+	// MaxTraceSize is the maximum number of spans allowed per trace.
 	// If a trace has more spans than this, it will be truncated and a warning will be added
 	MaxTraceSize int `mapstructure:"max_trace_size" valid:"optional"`
 	// EnableTracing determines whether traces will be emitted by jaeger-query.
diff --git a/cmd/jaeger/internal/extension/jaegerquery/querysvc/service.go b/cmd/jaeger/internal/extension/jaegerquery/querysvc/service.go
index a76ea1b719e..4a040336b2e 100644
--- a/cmd/jaeger/internal/extension/jaegerquery/querysvc/service.go
+++ b/cmd/jaeger/internal/extension/jaegerquery/querysvc/service.go
@@ -30,7 +30,7 @@ type QueryServiceOptions struct {
 	ArchiveTraceWriter tracestore.Writer
 	// MaxClockSkewAdjust is the maximum duration by which to adjust a span.
 	MaxClockSkewAdjust time.Duration
-	// MaxTraceSize is the max no. of spans allowed per trace.
+	// MaxTraceSize is the maximum number of spans allowed per trace.
 	// If a trace has more spans than this, it will be truncated and a warning will be added
 	MaxTraceSize int
 }

From 527fa523a00dd61d159c1362401347fc7ae26f6d Mon Sep 17 00:00:00 2001
From: Yuri Shkuro <github@ysh.us>
Date: Thu, 26 Feb 2026 12:36:58 -0500
Subject: [PATCH 17/22] fix

Signed-off-by: Yuri Shkuro <github@ysh.us>
---
 internal/jptrace/aggregator.go | 8 ++++++--
 internal/jptrace/spaniter.go   | 8 ++++++++
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/internal/jptrace/aggregator.go b/internal/jptrace/aggregator.go
index 1c93bc70d74..2c92e6002d2 100644
--- a/internal/jptrace/aggregator.go
+++ b/internal/jptrace/aggregator.go
@@ -18,6 +18,11 @@ func AggregateTraces(tracesSeq iter.Seq2[[]ptrace.Traces, error]) iter.Seq2[ptra
 	return AggregateTracesWithLimit(tracesSeq, 0)
 }
 
+// AggregateTracesWithLimit aggregates a sequence of trace batches into individual traces
+// but limits each trace size to maxSize spans. If maxSize is 0 or negative, there is no
+// limit and all spans will be included in the aggregated trace.
+//
+// The `tracesSeq` input must adhere to the chunking requirements of tracestore.Reader.GetTraces.
 func AggregateTracesWithLimit(tracesSeq iter.Seq2[[]ptrace.Traces, error], maxSize int) iter.Seq2[ptrace.Traces, error] {
 	return func(yield func(trace ptrace.Traces, err error) bool) {
 		currentTrace := ptrace.NewTraces()
@@ -34,8 +39,7 @@ func AggregateTracesWithLimit(tracesSeq iter.Seq2[[]ptrace.Traces, error], maxSi
 				if trace.SpanCount() == 0 {
 					continue
 				}
-				resources := trace.ResourceSpans()
-				traceID := resources.At(0).ScopeSpans().At(0).Spans().At(0).TraceID()
+				traceID := GetTraceID(trace)
 				if currentTraceID == traceID {
 					mergeTracesWithLimit(currentTrace, trace, maxSize, &spanCount)
 				} else {
diff --git a/internal/jptrace/spaniter.go b/internal/jptrace/spaniter.go
index 5ea70b8779c..78ea7b4baa4 100644
--- a/internal/jptrace/spaniter.go
+++ b/internal/jptrace/spaniter.go
@@ -6,6 +6,7 @@ package jptrace
 import (
 	"iter"
 
+	"go.opentelemetry.io/collector/pdata/pcommon"
 	"go.opentelemetry.io/collector/pdata/ptrace"
 )
 
@@ -38,3 +39,10 @@ func SpanIter(traces ptrace.Traces) iter.Seq2[SpanIterPos, ptrace.Span] {
 		}
 	}
 }
+
+func GetTraceID(traces ptrace.Traces) pcommon.TraceID {
+	for _, span := range SpanIter(traces) {
+		return span.TraceID()
+	}
+	return pcommon.NewTraceIDEmpty()
+}

From bf1cbf013b58e85abef4b2ee68bd8efb318b6fd6 Mon Sep 17 00:00:00 2001
From: Yuri Shkuro <github@ysh.us>
Date: Thu, 26 Feb 2026 12:38:40 -0500
Subject: [PATCH 18/22] fix

Signed-off-by: Yuri Shkuro <github@ysh.us>
---
 internal/jptrace/aggregator.go | 20 ++++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/internal/jptrace/aggregator.go b/internal/jptrace/aggregator.go
index 2c92e6002d2..95b6f8156ec 100644
--- a/internal/jptrace/aggregator.go
+++ b/internal/jptrace/aggregator.go
@@ -41,7 +41,7 @@ func AggregateTracesWithLimit(tracesSeq iter.Seq2[[]ptrace.Traces, error], maxSi
 				}
 				traceID := GetTraceID(trace)
 				if currentTraceID == traceID {
-					mergeTracesWithLimit(currentTrace, trace, maxSize, &spanCount)
+					spanCount = mergeTracesWithLimit(currentTrace, trace, maxSize, spanCount)
 				} else {
 					if currentTrace.SpanCount() > 0 {
 						if !yield(currentTrace, nil) {
@@ -69,29 +69,29 @@ func AggregateTracesWithLimit(tracesSeq iter.Seq2[[]ptrace.Traces, error], maxSi
 	}
 }
 
-func mergeTracesWithLimit(dest, src ptrace.Traces, maxSize int, spanCount *int) bool {
+func mergeTracesWithLimit(dest, src ptrace.Traces, maxSize int, spanCount int) int {
 	// early exit if already at max
-	if maxSize > 0 && *spanCount >= maxSize {
+	if maxSize > 0 && spanCount >= maxSize {
 		markTraceTruncated(dest, maxSize)
-		return true
+		return spanCount
 	}
 
 	incomingCount := src.SpanCount()
 	// check if we can merge all spans without exceeding limit
-	if maxSize <= 0 || *spanCount+incomingCount <= maxSize {
+	if maxSize <= 0 || spanCount+incomingCount <= maxSize {
 		MergeTraces(dest, src)
-		*spanCount += incomingCount
-		return false
+		spanCount += incomingCount
+		return spanCount
 	}
 
 	// partial copy
-	remaining := maxSize - *spanCount
+	remaining := maxSize - spanCount
 	if remaining > 0 {
 		copySpansUpToLimit(dest, src, remaining)
-		*spanCount = maxSize
+		spanCount = maxSize
 	}
 	markTraceTruncated(dest, maxSize)
-	return true
+	return spanCount
 }
 
 func copySpansUpToLimit(dest, src ptrace.Traces, limit int) {

From 9ad7aee63dcdd42d00a371134ecbe722b547b9b0 Mon Sep 17 00:00:00 2001
From: Yuri Shkuro <github@ysh.us>
Date: Thu, 26 Feb 2026 12:46:37 -0500
Subject: [PATCH 19/22] fix

Signed-off-by: Yuri Shkuro <github@ysh.us>
---
 internal/jptrace/aggregator.go      | 39 +++++++++++++-------
 internal/jptrace/aggregator_test.go | 57 +++++++++++++++++++++++++++++
 2 files changed, 82 insertions(+), 14 deletions(-)

diff --git a/internal/jptrace/aggregator.go b/internal/jptrace/aggregator.go
index 95b6f8156ec..925e78e79ce 100644
--- a/internal/jptrace/aggregator.go
+++ b/internal/jptrace/aggregator.go
@@ -49,15 +49,15 @@ func AggregateTracesWithLimit(tracesSeq iter.Seq2[[]ptrace.Traces, error], maxSi
 							return false
 						}
 					}
-					currentTrace = trace
 					currentTraceID = traceID
-					spanCount = trace.SpanCount()
-					if maxSize > 0 && spanCount > maxSize {
+					if maxSize > 0 && trace.SpanCount() > maxSize {
 						currentTrace = ptrace.NewTraces()
-						spanCount = 0
 						copySpansUpToLimit(currentTrace, trace, maxSize)
 						spanCount = maxSize
 						markTraceTruncated(currentTrace, maxSize)
+					} else {
+						currentTrace = trace
+						spanCount = trace.SpanCount()
 					}
 				}
 			}
@@ -70,9 +70,8 @@ func AggregateTracesWithLimit(tracesSeq iter.Seq2[[]ptrace.Traces, error], maxSi
 }
 
 func mergeTracesWithLimit(dest, src ptrace.Traces, maxSize int, spanCount int) int {
-	// early exit if already at max
+	// early exit if already at max; trace was already marked truncated when the limit was first hit
 	if maxSize > 0 && spanCount >= maxSize {
-		markTraceTruncated(dest, maxSize)
 		return spanCount
 	}
 
@@ -101,21 +100,33 @@ func copySpansUpToLimit(dest, src ptrace.Traces, limit int) {
 		if copied >= limit {
 			return
 		}
-		destResource := dest.ResourceSpans().AppendEmpty()
-		srcResource.Resource().CopyTo(destResource.Resource())
-		destResource.SetSchemaUrl(srcResource.SchemaUrl())
+		var destResource ptrace.ResourceSpans
+		resourceAdded := false
 
 		for _, srcScope := range srcResource.ScopeSpans().All() {
 			if copied >= limit {
-				return
+				break
 			}
-			destScope := destResource.ScopeSpans().AppendEmpty()
-			srcScope.Scope().CopyTo(destScope.Scope())
-			destScope.SetSchemaUrl(srcScope.SchemaUrl())
+			var destScope ptrace.ScopeSpans
+			scopeAdded := false
 
 			for _, span := range srcScope.Spans().All() {
 				if copied >= limit {
-					return
+					break
+				}
+				// Lazily create resource and scope containers only when a span is actually copied,
+				// to avoid leaving empty container artifacts in dest.
+				if !resourceAdded {
+					destResource = dest.ResourceSpans().AppendEmpty()
+					srcResource.Resource().CopyTo(destResource.Resource())
+					destResource.SetSchemaUrl(srcResource.SchemaUrl())
+					resourceAdded = true
+				}
+				if !scopeAdded {
+					destScope = destResource.ScopeSpans().AppendEmpty()
+					srcScope.Scope().CopyTo(destScope.Scope())
+					destScope.SetSchemaUrl(srcScope.SchemaUrl())
+					scopeAdded = true
 				}
 				span.CopyTo(destScope.Spans().AppendEmpty())
 				copied++
diff --git a/internal/jptrace/aggregator_test.go b/internal/jptrace/aggregator_test.go
index 327d58dc59f..48ecb011eda 100644
--- a/internal/jptrace/aggregator_test.go
+++ b/internal/jptrace/aggregator_test.go
@@ -199,6 +199,63 @@ func TestCopySpansUpToLimit(t *testing.T) {
 	assert.Equal(t, 3, dest.SpanCount())
 }
 
+func TestCopySpansUpToLimit_NoEmptyContainers(t *testing.T) {
+	// src has two resources: the first has no scopes, the second has spans.
+	// copySpansUpToLimit should not create an empty ResourceSpans for the first resource.
+	src := ptrace.NewTraces()
+	src.ResourceSpans().AppendEmpty() // empty resource, no scopes
+	spans := src.ResourceSpans().AppendEmpty().ScopeSpans().AppendEmpty().Spans()
+	for i := 0; i < 3; i++ {
+		spans.AppendEmpty().SetName("span")
+	}
+
+	dest := ptrace.NewTraces()
+	copySpansUpToLimit(dest, src, 2)
+
+	assert.Equal(t, 2, dest.SpanCount())
+	assert.Equal(t, 1, dest.ResourceSpans().Len(), "empty resource should not be copied")
+}
+
+func TestAggregateTracesWithLimit_MultiBatch(t *testing.T) {
+	// A trace that arrives in three batches should produce exactly one truncation
+	// warning even when subsequent batches arrive after the limit is already reached.
+	createBatch := func(traceID byte, spanCount int) ptrace.Traces {
+		trace := ptrace.NewTraces()
+		spans := trace.ResourceSpans().AppendEmpty().ScopeSpans().AppendEmpty().Spans()
+		for i := 0; i < spanCount; i++ {
+			span := spans.AppendEmpty()
+			span.SetTraceID(pcommon.TraceID([16]byte{traceID}))
+		}
+		return trace
+	}
+
+	// Limit is 3. Batch 1: 2 spans (under limit). Batch 2: 2 spans (partial copy, hits limit).
+	// Batch 3: 2 spans (already at limit, ignored).
+	tracesSeq := func(yield func([]ptrace.Traces, error) bool) {
+		if !yield([]ptrace.Traces{createBatch(1, 2)}, nil) {
+			return
+		}
+		if !yield([]ptrace.Traces{createBatch(1, 2)}, nil) {
+			return
+		}
+		yield([]ptrace.Traces{createBatch(1, 2)}, nil)
+	}
+
+	var result []ptrace.Traces
+	AggregateTracesWithLimit(tracesSeq, 3)(func(trace ptrace.Traces, _ error) bool {
+		result = append(result, trace)
+		return true
+	})
+
+	require.Len(t, result, 1)
+	assert.Equal(t, 3, result[0].SpanCount())
+
+	firstSpan := result[0].ResourceSpans().At(0).ScopeSpans().At(0).Spans().At(0)
+	warnings := GetWarnings(firstSpan)
+	assert.Len(t, warnings, 1, "should have exactly one truncation warning, not one per extra batch")
+	assert.Contains(t, warnings[0], fmt.Sprintf("trace has more than %d spans", 3))
+}
+
 func TestMarkAndCheckTruncated(t *testing.T) {
 	trace := ptrace.NewTraces()
 	firstSpan := trace.ResourceSpans().AppendEmpty().ScopeSpans().AppendEmpty().Spans().AppendEmpty()

From 43cf9a42e226ef9905aafd2946f949e25d84fc7d Mon Sep 17 00:00:00 2001
From: Yuri Shkuro <github@ysh.us>
Date: Thu, 26 Feb 2026 12:57:28 -0500
Subject: [PATCH 20/22] fix

Signed-off-by: Yuri Shkuro <github@ysh.us>
---
 internal/jptrace/aggregator.go | 23 ++++++++++++-----------
 1 file changed, 12 insertions(+), 11 deletions(-)

diff --git a/internal/jptrace/aggregator.go b/internal/jptrace/aggregator.go
index 925e78e79ce..ab1e6d7aacf 100644
--- a/internal/jptrace/aggregator.go
+++ b/internal/jptrace/aggregator.go
@@ -43,7 +43,7 @@ func AggregateTracesWithLimit(tracesSeq iter.Seq2[[]ptrace.Traces, error], maxSi
 				if currentTraceID == traceID {
 					spanCount = mergeTracesWithLimit(currentTrace, trace, maxSize, spanCount)
 				} else {
-					if currentTrace.SpanCount() > 0 {
+					if spanCount > 0 {
 						if !yield(currentTrace, nil) {
 							cont = false
 							return false
@@ -63,7 +63,7 @@ func AggregateTracesWithLimit(tracesSeq iter.Seq2[[]ptrace.Traces, error], maxSi
 			}
 			return true
 		})
-		if cont && currentTrace.SpanCount() > 0 {
+		if cont && spanCount > 0 {
 			yield(currentTrace, nil)
 		}
 	}
@@ -79,8 +79,7 @@ func mergeTracesWithLimit(dest, src ptrace.Traces, maxSize int, spanCount int) i
 	// check if we can merge all spans without exceeding limit
 	if maxSize <= 0 || spanCount+incomingCount <= maxSize {
 		MergeTraces(dest, src)
-		spanCount += incomingCount
-		return spanCount
+		return spanCount + incomingCount
 	}
 
 	// partial copy
@@ -105,14 +104,14 @@ func copySpansUpToLimit(dest, src ptrace.Traces, limit int) {
 
 		for _, srcScope := range srcResource.ScopeSpans().All() {
 			if copied >= limit {
-				break
+				return
 			}
 			var destScope ptrace.ScopeSpans
 			scopeAdded := false
 
 			for _, span := range srcScope.Spans().All() {
 				if copied >= limit {
-					break
+					return
 				}
 				// Lazily create resource and scope containers only when a span is actually copied,
 				// to avoid leaving empty container artifacts in dest.
@@ -146,9 +145,11 @@ func MergeTraces(dest, src ptrace.Traces) {
 }
 
 func markTraceTruncated(trace ptrace.Traces, maxSize int) {
-	SpanIter(trace)(func(_ SpanIterPos, span ptrace.Span) bool {
-		AddWarnings(span,
-			fmt.Sprintf("trace has more than %d spans, showing first %d spans only", maxSize, maxSize))
-		return false // stop after first span
-	})
+	for _, span := range SpanIter(trace) {
+		AddWarnings(
+			span,
+			fmt.Sprintf("trace has more than %d spans, showing first %d spans only", maxSize, maxSize),
+		)
+		return // stop after first span
+	}
 }

From 9b0b67eabee2b23efb48a0df36d69b5c5a2e7a61 Mon Sep 17 00:00:00 2001
From: Yuri Shkuro <github@ysh.us>
Date: Thu, 26 Feb 2026 13:15:06 -0500
Subject: [PATCH 21/22] fix

Signed-off-by: Yuri Shkuro <github@ysh.us>
---
 internal/jptrace/aggregator.go | 42 ++++++++++++++++++++--------------
 1 file changed, 25 insertions(+), 17 deletions(-)

diff --git a/internal/jptrace/aggregator.go b/internal/jptrace/aggregator.go
index ab1e6d7aacf..213e7b71e2e 100644
--- a/internal/jptrace/aggregator.go
+++ b/internal/jptrace/aggregator.go
@@ -27,7 +27,7 @@ func AggregateTracesWithLimit(tracesSeq iter.Seq2[[]ptrace.Traces, error], maxSi
 	return func(yield func(trace ptrace.Traces, err error) bool) {
 		currentTrace := ptrace.NewTraces()
 		currentTraceID := pcommon.NewTraceIDEmpty()
-		spanCount := 0
+		currentSpanCount := 0
 		cont := true
 
 		tracesSeq(func(traces []ptrace.Traces, err error) bool {
@@ -36,60 +36,68 @@ func AggregateTracesWithLimit(tracesSeq iter.Seq2[[]ptrace.Traces, error], maxSi
 				return false
 			}
 			for _, trace := range traces {
-				if trace.SpanCount() == 0 {
+				incomingSpanCount := trace.SpanCount()
+				if incomingSpanCount == 0 {
 					continue
 				}
 				traceID := GetTraceID(trace)
 				if currentTraceID == traceID {
-					spanCount = mergeTracesWithLimit(currentTrace, trace, maxSize, spanCount)
+					// same trace as current, merge it into the current trace, respecting the maxSize limit
+					currentSpanCount = mergeTracesWithLimit(currentTrace, currentSpanCount, trace, incomingSpanCount, maxSize)
 				} else {
-					if spanCount > 0 {
+					if currentSpanCount > 0 {
 						if !yield(currentTrace, nil) {
 							cont = false
 							return false
 						}
 					}
 					currentTraceID = traceID
-					if maxSize > 0 && trace.SpanCount() > maxSize {
+					if maxSize > 0 && incomingSpanCount > maxSize {
 						currentTrace = ptrace.NewTraces()
 						copySpansUpToLimit(currentTrace, trace, maxSize)
-						spanCount = maxSize
+						currentSpanCount = maxSize
 						markTraceTruncated(currentTrace, maxSize)
 					} else {
+						// Optimization: when incoming trace fits within the limit (or there is no limit),
+						// we can skip the copy and use it directly as the current trace.
 						currentTrace = trace
-						spanCount = trace.SpanCount()
+						currentSpanCount = incomingSpanCount
 					}
 				}
 			}
 			return true
 		})
-		if cont && spanCount > 0 {
+		// Emit the last accumulated trace if non-empty.
+		// `cont` guards against calling yield after consumer already returned false.
+		if cont && currentSpanCount > 0 {
 			yield(currentTrace, nil)
 		}
 	}
 }
 
-func mergeTracesWithLimit(dest, src ptrace.Traces, maxSize int, spanCount int) int {
+// mergeTracesWithLimit merges src into dest, respecting the maxSize span limit.
+// destCount and srcCount are the pre-computed span counts for dest and src respectively.
+// Returns the updated span count of dest after the merge.
+// If maxSize <= 0, all spans are merged without limit.
+func mergeTracesWithLimit(dest ptrace.Traces, destCount int, src ptrace.Traces, srcCount int, maxSize int) int {
 	// early exit if already at max; trace was already marked truncated when the limit was first hit
-	if maxSize > 0 && spanCount >= maxSize {
-		return spanCount
+	if maxSize > 0 && destCount >= maxSize {
+		return destCount
 	}
 
-	incomingCount := src.SpanCount()
 	// check if we can merge all spans without exceeding limit
-	if maxSize <= 0 || spanCount+incomingCount <= maxSize {
+	if maxSize <= 0 || destCount+srcCount <= maxSize {
 		MergeTraces(dest, src)
-		return spanCount + incomingCount
+		return destCount + srcCount
 	}
 
 	// partial copy
-	remaining := maxSize - spanCount
+	remaining := maxSize - destCount
 	if remaining > 0 {
 		copySpansUpToLimit(dest, src, remaining)
-		spanCount = maxSize
 	}
 	markTraceTruncated(dest, maxSize)
-	return spanCount
+	return maxSize
 }
 
 func copySpansUpToLimit(dest, src ptrace.Traces, limit int) {

From 5ae83621bfc3f0c4b55bd8a113bc523e43fe01f3 Mon Sep 17 00:00:00 2001
From: Yuri Shkuro <github@ysh.us>
Date: Thu, 26 Feb 2026 13:34:51 -0500
Subject: [PATCH 22/22] fix

Signed-off-by: Yuri Shkuro <github@ysh.us>
---
 .../extension/jaegerquery/internal/flags.go   |  4 +-
 .../extension/jaegerquery/querysvc/service.go |  4 +-
 internal/jptrace/aggregator.go                | 32 ++++----
 internal/jptrace/aggregator_test.go           | 79 +++++++++++++++++++
 internal/jptrace/spaniter_test.go             | 16 ++++
 5 files changed, 117 insertions(+), 18 deletions(-)

diff --git a/cmd/jaeger/internal/extension/jaegerquery/internal/flags.go b/cmd/jaeger/internal/extension/jaegerquery/internal/flags.go
index acce6ff1eea..b4a3720d065 100644
--- a/cmd/jaeger/internal/extension/jaegerquery/internal/flags.go
+++ b/cmd/jaeger/internal/extension/jaegerquery/internal/flags.go
@@ -36,8 +36,8 @@ type QueryOptions struct {
 	Tenancy tenancy.Options `mapstructure:"multi_tenancy"`
 	// MaxClockSkewAdjust is the maximum duration by which jaeger-query will adjust a span.
 	MaxClockSkewAdjust time.Duration `mapstructure:"max_clock_skew_adjust"  valid:"optional"`
-	// MaxTraceSize is the maximum number of spans allowed per trace.
-	// If a trace has more spans than this, it will be truncated and a warning will be added
+	// MaxTraceSize is the maximum number of spans allowed per trace. A value of 0 (default) means unlimited.
+	// If a trace has more spans than this limit, it will be truncated and a warning will be added.
 	MaxTraceSize int `mapstructure:"max_trace_size" valid:"optional"`
 	// EnableTracing determines whether traces will be emitted by jaeger-query.
 	EnableTracing bool `mapstructure:"enable_tracing"`
diff --git a/cmd/jaeger/internal/extension/jaegerquery/querysvc/service.go b/cmd/jaeger/internal/extension/jaegerquery/querysvc/service.go
index 4a040336b2e..7198a2f7c08 100644
--- a/cmd/jaeger/internal/extension/jaegerquery/querysvc/service.go
+++ b/cmd/jaeger/internal/extension/jaegerquery/querysvc/service.go
@@ -30,8 +30,8 @@ type QueryServiceOptions struct {
 	ArchiveTraceWriter tracestore.Writer
 	// MaxClockSkewAdjust is the maximum duration by which to adjust a span.
 	MaxClockSkewAdjust time.Duration
-	// MaxTraceSize is the maximum number of spans allowed per trace.
-	// If a trace has more spans than this, it will be truncated and a warning will be added
+	// MaxTraceSize is the maximum number of spans allowed per trace. A value of 0 (default) means unlimited.
+	// If a trace has more spans than this limit, it will be truncated and a warning will be added.
 	MaxTraceSize int
 }
 
diff --git a/internal/jptrace/aggregator.go b/internal/jptrace/aggregator.go
index 213e7b71e2e..515bd4b52db 100644
--- a/internal/jptrace/aggregator.go
+++ b/internal/jptrace/aggregator.go
@@ -28,6 +28,7 @@ func AggregateTracesWithLimit(tracesSeq iter.Seq2[[]ptrace.Traces, error], maxSi
 		currentTrace := ptrace.NewTraces()
 		currentTraceID := pcommon.NewTraceIDEmpty()
 		currentSpanCount := 0
+		currentTruncated := false
 		cont := true
 
 		tracesSeq(func(traces []ptrace.Traces, err error) bool {
@@ -43,7 +44,12 @@ func AggregateTracesWithLimit(tracesSeq iter.Seq2[[]ptrace.Traces, error], maxSi
 				traceID := GetTraceID(trace)
 				if currentTraceID == traceID {
 					// same trace as current, merge it into the current trace, respecting the maxSize limit
-					currentSpanCount = mergeTracesWithLimit(currentTrace, currentSpanCount, trace, incomingSpanCount, maxSize)
+					var truncated bool
+					currentSpanCount, truncated = mergeTracesWithLimit(currentTrace, currentSpanCount, trace, incomingSpanCount, maxSize)
+					if truncated && !currentTruncated {
+						markTraceTruncated(currentTrace, maxSize)
+						currentTruncated = true
+					}
 				} else {
 					if currentSpanCount > 0 {
 						if !yield(currentTrace, nil) {
@@ -52,11 +58,13 @@ func AggregateTracesWithLimit(tracesSeq iter.Seq2[[]ptrace.Traces, error], maxSi
 						}
 					}
 					currentTraceID = traceID
+					currentTruncated = false
 					if maxSize > 0 && incomingSpanCount > maxSize {
 						currentTrace = ptrace.NewTraces()
 						copySpansUpToLimit(currentTrace, trace, maxSize)
 						currentSpanCount = maxSize
 						markTraceTruncated(currentTrace, maxSize)
+						currentTruncated = true
 					} else {
 						// Optimization: when incoming trace fits within the limit (or there is no limit),
 						// we can skip the copy and use it directly as the current trace.
@@ -77,27 +85,23 @@ func AggregateTracesWithLimit(tracesSeq iter.Seq2[[]ptrace.Traces, error], maxSi
 
 // mergeTracesWithLimit merges src into dest, respecting the maxSize span limit.
 // destCount and srcCount are the pre-computed span counts for dest and src respectively.
-// Returns the updated span count of dest after the merge.
-// If maxSize <= 0, all spans are merged without limit.
-func mergeTracesWithLimit(dest ptrace.Traces, destCount int, src ptrace.Traces, srcCount int, maxSize int) int {
-	// early exit if already at max; trace was already marked truncated when the limit was first hit
+// Returns the updated span count of dest and whether the trace was truncated (true if
+// src spans were dropped due to the limit). If maxSize <= 0, all spans are merged without limit.
+func mergeTracesWithLimit(dest ptrace.Traces, destCount int, src ptrace.Traces, srcCount int, maxSize int) (int, bool) {
+	// early exit if already at max
 	if maxSize > 0 && destCount >= maxSize {
-		return destCount
+		return destCount, true
 	}
 
 	// check if we can merge all spans without exceeding limit
 	if maxSize <= 0 || destCount+srcCount <= maxSize {
 		MergeTraces(dest, src)
-		return destCount + srcCount
+		return destCount + srcCount, false
 	}
 
-	// partial copy
-	remaining := maxSize - destCount
-	if remaining > 0 {
-		copySpansUpToLimit(dest, src, remaining)
-	}
-	markTraceTruncated(dest, maxSize)
-	return maxSize
+	// partial copy: only copy the spans that fit
+	copySpansUpToLimit(dest, src, maxSize-destCount)
+	return maxSize, true
 }
 
 func copySpansUpToLimit(dest, src ptrace.Traces, limit int) {
diff --git a/internal/jptrace/aggregator_test.go b/internal/jptrace/aggregator_test.go
index 48ecb011eda..3cc830ae863 100644
--- a/internal/jptrace/aggregator_test.go
+++ b/internal/jptrace/aggregator_test.go
@@ -199,6 +199,47 @@ func TestCopySpansUpToLimit(t *testing.T) {
 	assert.Equal(t, 3, dest.SpanCount())
 }
 
+func TestCopySpansUpToLimit_MultipleResourceSpans(t *testing.T) {
+	src := ptrace.NewTraces()
+	rs0 := src.ResourceSpans().AppendEmpty()
+	ss0 := rs0.ScopeSpans().AppendEmpty()
+	ss0.Spans().AppendEmpty().SetName("rs0-span0")
+	ss0.Spans().AppendEmpty().SetName("rs0-span1")
+	rs1 := src.ResourceSpans().AppendEmpty()
+	ss1 := rs1.ScopeSpans().AppendEmpty()
+	ss1.Spans().AppendEmpty().SetName("rs1-span0")
+	ss1.Spans().AppendEmpty().SetName("rs1-span1")
+
+	dest := ptrace.NewTraces()
+	copySpansUpToLimit(dest, src, 3)
+
+	require.Equal(t, 3, dest.SpanCount())
+	require.Equal(t, 2, dest.ResourceSpans().Len())
+	assert.Equal(t, 2, dest.ResourceSpans().At(0).ScopeSpans().At(0).Spans().Len())
+	assert.Equal(t, 1, dest.ResourceSpans().At(1).ScopeSpans().At(0).Spans().Len())
+}
+
+func TestCopySpansUpToLimit_MultipleScopeSpans(t *testing.T) {
+	src := ptrace.NewTraces()
+	rs := src.ResourceSpans().AppendEmpty()
+	ss0 := rs.ScopeSpans().AppendEmpty()
+	ss0.Spans().AppendEmpty().SetName("ss0-span0")
+	ss0.Spans().AppendEmpty().SetName("ss0-span1")
+	ss1 := rs.ScopeSpans().AppendEmpty()
+	ss1.Spans().AppendEmpty().SetName("ss1-span0")
+	ss1.Spans().AppendEmpty().SetName("ss1-span1")
+
+	dest := ptrace.NewTraces()
+	copySpansUpToLimit(dest, src, 3)
+
+	require.Equal(t, 3, dest.SpanCount())
+	require.Equal(t, 1, dest.ResourceSpans().Len())
+	destScopes := dest.ResourceSpans().At(0).ScopeSpans()
+	require.Equal(t, 2, destScopes.Len())
+	assert.Equal(t, 2, destScopes.At(0).Spans().Len())
+	assert.Equal(t, 1, destScopes.At(1).Spans().Len())
+}
+
 func TestCopySpansUpToLimit_NoEmptyContainers(t *testing.T) {
 	// src has two resources: the first has no scopes, the second has spans.
 	// copySpansUpToLimit should not create an empty ResourceSpans for the first resource.
@@ -256,6 +297,44 @@ func TestAggregateTracesWithLimit_MultiBatch(t *testing.T) {
 	assert.Contains(t, warnings[0], fmt.Sprintf("trace has more than %d spans", 3))
 }
 
+// TestAggregateTracesWithLimit_ExactLimitThenOverflow specifically tests the scenario
+// where the first batch fills the trace to exactly maxSize (no warning yet), and a
+// subsequent batch then causes the first overflow and must trigger the truncation warning.
+func TestAggregateTracesWithLimit_ExactLimitThenOverflow(t *testing.T) {
+	createBatch := func(traceID byte, spanCount int) ptrace.Traces {
+		trace := ptrace.NewTraces()
+		spans := trace.ResourceSpans().AppendEmpty().ScopeSpans().AppendEmpty().Spans()
+		for i := 0; i < spanCount; i++ {
+			span := spans.AppendEmpty()
+			span.SetTraceID(pcommon.TraceID([16]byte{traceID}))
+		}
+		return trace
+	}
+
+	// Batch 1 has exactly maxSize spans — fits without truncation, no warning added yet.
+	// Batch 2 has 1 more span — must be dropped AND must trigger the warning.
+	tracesSeq := func(yield func([]ptrace.Traces, error) bool) {
+		if !yield([]ptrace.Traces{createBatch(1, 3)}, nil) {
+			return
+		}
+		yield([]ptrace.Traces{createBatch(1, 1)}, nil)
+	}
+
+	var result []ptrace.Traces
+	AggregateTracesWithLimit(tracesSeq, 3)(func(trace ptrace.Traces, _ error) bool {
+		result = append(result, trace)
+		return true
+	})
+
+	require.Len(t, result, 1)
+	assert.Equal(t, 3, result[0].SpanCount())
+
+	firstSpan := result[0].ResourceSpans().At(0).ScopeSpans().At(0).Spans().At(0)
+	warnings := GetWarnings(firstSpan)
+	assert.Len(t, warnings, 1, "overflow after exact-limit batch must produce exactly one truncation warning")
+	assert.Contains(t, warnings[0], fmt.Sprintf("trace has more than %d spans", 3))
+}
+
 func TestMarkAndCheckTruncated(t *testing.T) {
 	trace := ptrace.NewTraces()
 	firstSpan := trace.ResourceSpans().AppendEmpty().ScopeSpans().AppendEmpty().Spans().AppendEmpty()
diff --git a/internal/jptrace/spaniter_test.go b/internal/jptrace/spaniter_test.go
index 77d0e8b762a..9fccfb9b0cd 100644
--- a/internal/jptrace/spaniter_test.go
+++ b/internal/jptrace/spaniter_test.go
@@ -7,6 +7,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"go.opentelemetry.io/collector/pdata/pcommon"
 	"go.opentelemetry.io/collector/pdata/ptrace"
 )
 
@@ -92,3 +93,18 @@ func TestSpanIterStopIteration(t *testing.T) {
 	assert.Len(t, spans, 1)
 	assert.Equal(t, "span-1", spans[0].Name())
 }
+
+func TestGetTraceID(t *testing.T) {
+	t.Run("empty traces returns empty TraceID", func(t *testing.T) {
+		traces := ptrace.NewTraces()
+		assert.Equal(t, pcommon.NewTraceIDEmpty(), GetTraceID(traces))
+	})
+
+	t.Run("returns TraceID of first span", func(t *testing.T) {
+		traces := ptrace.NewTraces()
+		span := traces.ResourceSpans().AppendEmpty().ScopeSpans().AppendEmpty().Spans().AppendEmpty()
+		traceID := pcommon.TraceID([16]byte{1, 2, 3})
+		span.SetTraceID(traceID)
+		assert.Equal(t, traceID, GetTraceID(traces))
+	})
+}