Merge remote-tracking branch 'origin/main' into issue_5601_spanlist

pkg/statementbuilder/aistatementbuilder/span_list_test.go

@@ -0,0 +1,155 @@
+package aistatementbuilder
+
+import (
+	"context"
+	"testing"
+
+	qbtypes "github.com/SigNoz/signoz/pkg/types/querybuildertypes/querybuildertypesv5"
+	"github.com/SigNoz/signoz/pkg/types/telemetrytypes"
+	"github.com/SigNoz/signoz/pkg/valuer"
+	"github.com/stretchr/testify/require"
+)
+
+// Span list with a mixed filter: gen_ai spans matching the span-level part, in
+// traces whose window-clipped aggregates satisfy the trace-level part (the
+// __trace_scope qualification on the delegated path).
+func TestBuild_FullSQL_SpanList_TraceScoped(t *testing.T) {
+	b := newTestBuilder(t)
+	stmt, err := b.Build(context.Background(), valuer.UUID{}, testStartMs, testEndMs, qbtypes.RequestTypeRaw,
+		qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
+			Signal: telemetrytypes.SignalTraces,
+			Filter: &qbtypes.Filter{Expression: "gen_ai.request.model = 'gpt-4o-mini' AND trace.output_tokens > 1000"},
+			Limit:  10,
+		}, nil)
+	require.NoError(t, err)
+
+	assertSQLEqual(t, `
+WITH __trace_scope AS (
+    SELECT trace_id,
+        sum(multiIf(mapContains(attributes_number, 'gen_ai.usage.output_tokens'), toFloat64(attributes_number['gen_ai.usage.output_tokens']), NULL)) AS output_tokens
+    FROM signoz_traces.distributed_signoz_index_v3
+    WHERE timestamp >= '1747947419000000000'
+      AND timestamp < '1747983448000000000'
+      AND ts_bucket_start >= 1747945619
+      AND ts_bucket_start <= 1747983448
+      AND ((mapContains(attributes_string, 'gen_ai.request.model') OR mapContains(attributes_string, 'gen_ai.tool.name') OR mapContains(attributes_string, 'gen_ai.agent.name')))
+    GROUP BY trace_id
+    HAVING output_tokens > 1000
+)
+SELECT timestamp AS __SELECT_KEY_0_timestamp, trace_id AS __SELECT_KEY_1_trace_id, span_id AS __SELECT_KEY_2_span_id,
+    trace_state AS __SELECT_KEY_3_trace_state, parent_span_id AS __SELECT_KEY_4_parent_span_id, flags AS __SELECT_KEY_5_flags,
+    name AS __SELECT_KEY_6_name, kind AS __SELECT_KEY_7_kind, kind_string AS __SELECT_KEY_8_kind_string, duration_nano AS __SELECT_KEY_9_duration_nano,
+    status_code AS __SELECT_KEY_10_status_code, status_message AS __SELECT_KEY_11_status_message,
+    status_code_string AS __SELECT_KEY_12_status_code_string, events AS __SELECT_KEY_13_events, links AS __SELECT_KEY_14_links,
+    response_status_code AS __SELECT_KEY_15_response_status_code, external_http_url AS __SELECT_KEY_16_external_http_url,
+    http_url AS __SELECT_KEY_17_http_url, external_http_method AS __SELECT_KEY_18_external_http_method,
+    http_method AS __SELECT_KEY_19_http_method, http_host AS __SELECT_KEY_20_http_host, db_name AS __SELECT_KEY_21_db_name,
+    db_operation AS __SELECT_KEY_22_db_operation, has_error AS __SELECT_KEY_23_has_error, is_remote AS __SELECT_KEY_24_is_remote,
+    attributes_string, attributes_number, attributes_bool, resources_string
+FROM signoz_traces.distributed_signoz_index_v3
+WHERE trace_id GLOBAL IN (SELECT trace_id FROM __trace_scope)
+  AND (((mapContains(attributes_string, 'gen_ai.request.model')
+        OR mapContains(attributes_string, 'gen_ai.tool.name')
+        OR mapContains(attributes_string, 'gen_ai.agent.name')))
+    AND ((attributes_string['gen_ai.request.model'] = 'gpt-4o-mini'
+        AND mapContains(attributes_string, 'gen_ai.request.model'))))
+  AND timestamp >= '1747947419000000000'
+  AND timestamp < '1747983448000000000'
+  AND ts_bucket_start >= 1747945619
+  AND ts_bucket_start <= 1747983448
+LIMIT 10
+`, stmt)
+}
+
+// A raw filter mixing a resource attribute with a trace-level condition: the resource
+// part flows through the delegate's fingerprint machinery (__resource_filter CTE),
+// the trace-level part becomes the __trace_scope qualification.
+func TestBuild_SpanList_ResourcePlusTraceFilter(t *testing.T) {
+	b := newTestBuilder(t)
+	stmt, err := b.Build(context.Background(), valuer.UUID{}, testStartMs, testEndMs, qbtypes.RequestTypeRaw,
+		qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
+			Signal: telemetrytypes.SignalTraces,
+			Filter: &qbtypes.Filter{Expression: "resource.service.name = 'checkout' AND trace.output_tokens > 1000"},
+			Limit:  10,
+		}, nil)
+	require.NoError(t, err)
+
+	got := renderSQL(t, stmt)
+	require.Contains(t, got, "__resource_filter AS (")
+	require.Contains(t, got, "resource_fingerprint GLOBAL IN (SELECT fingerprint FROM __resource_filter)")
+	require.Contains(t, got, "__trace_scope AS (")
+	require.Contains(t, got, "trace_id GLOBAL IN (SELECT trace_id FROM __trace_scope)")
+	require.Contains(t, got, "HAVING output_tokens > 1000")
+}
+
+// Without a trace-level condition nothing changes: the span list stays a single
+// gated span scan (no __trace_scope CTE).
+func TestBuild_SpanList_NoTraceFilter_NoScope(t *testing.T) {
+	b := newTestBuilder(t)
+	stmt, err := b.Build(context.Background(), valuer.UUID{}, testStartMs, testEndMs, qbtypes.RequestTypeRaw,
+		qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
+			Signal: telemetrytypes.SignalTraces,
+			Filter: &qbtypes.Filter{Expression: "gen_ai.request.model = 'gpt-4o-mini'"},
+			Limit:  10,
+		}, nil)
+	require.NoError(t, err)
+	require.NotContains(t, stmt.Query, "__trace_scope")
+}
+
+// The span-list trace-level filter shares the trace list's rules: output-only
+// aggregates are rejected, OR-mixing the two classes is rejected, and explicitly
+// trace-level order keys get a targeted error — while bare span columns that happen
+// to share a name with an aggregate alias (duration_nano) stay orderable.
+func TestBuild_SpanList_TraceFilter_Validation(t *testing.T) {
+	b := newTestBuilder(t)
+	build := func(q qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]) error {
+		q.Signal = telemetrytypes.SignalTraces
+		_, err := b.Build(context.Background(), valuer.UUID{}, testStartMs, testEndMs, qbtypes.RequestTypeRaw, q, nil)
+		return err
+	}
+
+	err := build(qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
+		Filter: &qbtypes.Filter{Expression: "trace.span_count > 3"},
+	})
+	require.ErrorContains(t, err, `aggregate "span_count" cannot be used`)
+
+	err = build(qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
+		Filter: &qbtypes.Filter{Expression: "trace.output_tokens > 1000 OR kind_string = 'Client'"},
+	})
+	require.ErrorContains(t, err, "cannot be combined")
+
+	err = build(qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
+		Order: []qbtypes.OrderBy{{Key: qbtypes.OrderByKey{TelemetryFieldKey: telemetrytypes.TelemetryFieldKey{Name: "trace.output_tokens"}}}},
+	})
+	require.ErrorContains(t, err, `ordering the span list by trace-level aggregate "trace.output_tokens" is not supported`)
+
+	err = build(qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
+		Order: []qbtypes.OrderBy{{Key: qbtypes.OrderByKey{TelemetryFieldKey: telemetrytypes.TelemetryFieldKey{Name: "duration_nano"}}, Direction: qbtypes.OrderDirectionDesc}},
+		Limit: 10,
+	})
+	require.NoError(t, err, "bare duration_nano is a span column, not a trace-level key")
+}
+
+// Variables in a trace-level condition on the span list get the trace list's
+// treatment: substituted as literals, __all__ drops the condition (no scope CTE).
+func TestBuild_SpanList_TraceFilter_Variables(t *testing.T) {
+	b := newTestBuilder(t)
+	build := func(expr string, vars map[string]qbtypes.VariableItem) (*qbtypes.Statement, error) {
+		return b.Build(context.Background(), valuer.UUID{}, testStartMs, testEndMs, qbtypes.RequestTypeRaw,
+			qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
+				Signal: telemetrytypes.SignalTraces,
+				Filter: &qbtypes.Filter{Expression: expr},
+				Limit:  10,
+			}, vars)
+	}
+
+	stmt, err := build("trace.output_tokens > $threshold",
+		map[string]qbtypes.VariableItem{"threshold": {Value: 700}})
+	require.NoError(t, err)
+	require.Contains(t, stmt.Query, "HAVING output_tokens > 700")
+
+	stmt, err = build("trace.output_tokens > $threshold",
+		map[string]qbtypes.VariableItem{"threshold": {Type: qbtypes.DynamicVariableType, Value: "__all__"}})
+	require.NoError(t, err)
+	require.NotContains(t, stmt.Query, "__trace_scope")
+}

pkg/statementbuilder/scopedtracesstatementbuilder/statement_builder.go

@@ -31,15 +31,24 @@ var (
 // (e.g. gen_ai spans); the TraceScope decides which spans are in scope and which
 // per-trace columns to compute.
 type scopedTraceStatementBuilder struct {
-	logger                    *slog.Logger
-	metadataStore             telemetrytypes.MetadataStore
-	fm                        qbtypes.FieldMapper
-	cb                        qbtypes.ConditionBuilder
-	scope                     TraceScope
-	traceStmtBuilder          qbtypes.StatementBuilder[qbtypes.TraceAggregation]
+	logger           *slog.Logger
+	metadataStore    telemetrytypes.MetadataStore
+	fm               qbtypes.FieldMapper
+	cb               qbtypes.ConditionBuilder
+	scope            TraceScope
+	traceStmtBuilder qbtypes.StatementBuilder[qbtypes.TraceAggregation]
+	// scopedDelegate is traceStmtBuilder's trace-scoping capability, resolved once at
+	// construction; nil when the delegate cannot constrain a query by trace ids.
+	scopedDelegate            traceScopedStatementBuilder
 	resourceFilterStmtBuilder qbtypes.StatementBuilder[qbtypes.TraceAggregation]
 }
 
+// traceScopedStatementBuilder is the delegate's optional capability of constraining a
+// query to a set of trace ids (implemented by the traces statement builder).
+type traceScopedStatementBuilder interface {
+	BuildTraceScoped(ctx context.Context, orgID valuer.UUID, start, end uint64, requestType qbtypes.RequestType, query qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation], variables map[string]qbtypes.VariableItem, traceScope *qbtypes.Statement) (*qbtypes.Statement, error)
+}
+
 var _ qbtypes.StatementBuilder[qbtypes.TraceAggregation] = (*scopedTraceStatementBuilder)(nil)
 
 // NewFactory returns a provider factory for a scoped trace statement builder. The
@@ -90,6 +99,7 @@ func NewScopedTraceStatementBuilder(
 		fl,
 	)
 
+	scopedDelegate, _ := traceStmtBuilder.(traceScopedStatementBuilder)
 	return &scopedTraceStatementBuilder{
 		logger:                    scopedSettings.Logger(),
 		metadataStore:             metadataStore,
@@ -97,6 +107,7 @@ func NewScopedTraceStatementBuilder(
 		cb:                        conditionBuilder,
 		scope:                     scope,
 		traceStmtBuilder:          traceStmtBuilder,
+		scopedDelegate:            scopedDelegate,
 		resourceFilterStmtBuilder: resourceFilterStmtBuilder,
 	}
 }
@@ -114,14 +125,18 @@ func (b *scopedTraceStatementBuilder) Build(
 	case qbtypes.RequestTypeTrace:
 		return b.buildTraceListQuery(ctx, orgID, querybuilder.ToNanoSecs(start), querybuilder.ToNanoSecs(end), query, variables)
 	case qbtypes.RequestTypeRaw:
+		if err := b.validateRawOrderKeys(query); err != nil {
+			return nil, err
+		}
 		return b.buildDelegated(ctx, orgID, start, end, requestType, query, variables)
 	default:
 		return nil, ErrUnsupportedRequestType
 	}
 }
 
-// buildDelegated ANDs the base gate into the user filter and delegates to the
-// standard trace builder (the span-list / raw path).
+// buildDelegated serves the span-list / raw path: the gate is ANDed into the filter's
+// span-level part and the query delegates to the standard trace builder; a trace-level
+// part becomes a qualification the delegate constrains trace_id by.
 func (b *scopedTraceStatementBuilder) buildDelegated(
 	ctx context.Context,
 	orgID valuer.UUID,
@@ -130,17 +145,127 @@ func (b *scopedTraceStatementBuilder) buildDelegated(
 	query qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation],
 	variables map[string]qbtypes.VariableItem,
 ) (*qbtypes.Statement, error) {
+	var spanExpr, traceExpr string
+	if query.Filter != nil && strings.TrimSpace(query.Filter.Expression) != "" {
+		var err error
+		spanExpr, traceExpr, err = querybuilder.SplitFilterForAggregates(query.Filter.Expression, b.aggregateAliasSet())
+		if err != nil {
+			return nil, err
+		}
+	}
+	traceExpr, err := substituteTraceLevelVariables(traceExpr, variables)
+	if err != nil {
+		return nil, err
+	}
+
 	gate := b.scope.FilterExpression
 	expr := gate
-	if query.Filter != nil && strings.TrimSpace(query.Filter.Expression) != "" {
-		expr = fmt.Sprintf("(%s) AND (%s)", gate, query.Filter.Expression)
+	if strings.TrimSpace(spanExpr) != "" {
+		expr = fmt.Sprintf("(%s) AND (%s)", gate, spanExpr)
 	}
 
 	// shallow copy; only Filter is replaced, caller's query untouched
 	gated := query
 	gated.Filter = &qbtypes.Filter{Expression: expr}
 
-	return b.traceStmtBuilder.Build(ctx, orgID, start, end, requestType, gated, variables)
+	if strings.TrimSpace(traceExpr) == "" {
+		return b.traceStmtBuilder.Build(ctx, orgID, start, end, requestType, gated, variables)
+	}
+
+	if b.scopedDelegate == nil {
+		return nil, errors.NewInternalf(errors.CodeInternal, "trace statement builder does not support trace-scoped queries")
+	}
+	scope, err := b.buildTraceScopeStatement(ctx, orgID, querybuilder.ToNanoSecs(start), querybuilder.ToNanoSecs(end), traceExpr)
+	if err != nil {
+		return nil, err
+	}
+	return b.scopedDelegate.BuildTraceScoped(ctx, orgID, start, end, requestType, gated, variables, scope)
+}
+
+// validateRawOrderKeys rejects ordering the span list by an explicitly trace-level
+// aggregate (trace. prefix or trace field context) — the per-trace value does not
+// exist on span rows. Bare names pass through: they may legitimately be span columns
+// (duration_nano, timestamp).
+func (b *scopedTraceStatementBuilder) validateRawOrderKeys(query qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]) error {
+	aliases := b.aggregateAliasSet()
+	for _, o := range query.Order {
+		key := telemetrytypes.GetFieldKeyFromKeyText(o.Key.Name)
+		if _, ok := aliases[key.Name]; !ok {
+			continue
+		}
+		if key.FieldContext == telemetrytypes.FieldContextTrace || o.Key.FieldContext == telemetrytypes.FieldContextTrace {
+			return errors.NewInvalidInputf(errors.CodeInvalidInput,
+				"ordering the span list by trace-level aggregate %q is not supported; order by span columns instead (e.g. timestamp, duration_nano)", o.Key.Name)
+		}
+	}
+	return nil
+}
+
+// substituteTraceLevelVariables resolves query variables in a trace-level expression.
+// The span-level parts bind variables via PrepareWhereClause; trace-level parts are
+// text rewrites, so variables become literals (a dynamic __all__ drops its condition).
+func substituteTraceLevelVariables(expr string, variables map[string]qbtypes.VariableItem) (string, error) {
+	if strings.TrimSpace(expr) == "" || len(variables) == 0 {
+		return expr, nil
+	}
+	return qbvariables.ReplaceVariablesInExpression(expr, variables)
+}
+
+// buildTraceScopeStatement builds the __trace_scope statement: trace ids whose
+// window-clipped per-trace aggregates satisfy traceExpr. The same scan as the matched
+// CTE, minus its span-filter widening, resource prune, ordering and pagination.
+// start/end are ns.
+func (b *scopedTraceStatementBuilder) buildTraceScopeStatement(ctx context.Context, orgID valuer.UUID, start, end uint64, traceExpr string) (*qbtypes.Statement, error) {
+	keys, err := b.fetchKeys(ctx, orgID)
+	if err != nil {
+		return nil, err
+	}
+	sb := sqlbuilder.NewSelectBuilder()
+	maskExpr, resolved, err := b.resolveFor(ctx, orgID, start, end, keys, sb)
+	if err != nil {
+		return nil, err
+	}
+	orderableSet := orderableAliasSet(resolved)
+	if err := validateAggregateFilter(traceExpr, orderableSet); err != nil {
+		return nil, err
+	}
+
+	needed := neededMatchedAliases(nil, traceExpr, orderableSet)
+	selects := []string{"trace_id"}
+	for _, rc := range resolved {
+		if _, ok := needed[rc.alias]; !ok {
+			continue
+		}
+		selects = append(selects, rc.expr+" AS "+quoteAlias(rc.alias))
+	}
+	sb.Select(selects...)
+	sb.From(fmt.Sprintf("%s.%s", tracestelemetryschema.DBName, tracestelemetryschema.SpanIndexV3TableName))
+	sb.Where(
+		sb.GE("timestamp", fmt.Sprintf("%d", start)),
+		sb.L("timestamp", fmt.Sprintf("%d", end)),
+		sb.GE("ts_bucket_start", start/querybuilder.NsToSeconds-querybuilder.BucketAdjustment),
+		sb.LE("ts_bucket_start", end/querybuilder.NsToSeconds),
+		"("+maskExpr+")",
+	)
+	sb.GroupBy("trace_id")
+
+	// the rewriter matches raw key text, so map the trace. form alongside the bare name
+	columnMap := make(map[string]string, len(orderableSet)*2)
+	for a := range orderableSet {
+		columnMap[a] = quoteAlias(a)
+		columnMap[telemetrytypes.FieldContextTrace.StringValue()+"."+a] = quoteAlias(a)
+	}
+	hv, err := querybuilder.NewHavingExpressionRewriter().Rewrite(traceExpr, columnMap)
+	if err != nil {
+		return nil, err
+	}
+	if hv != "" {
+		// escape user text so a literal $ isn't read as an arg marker
+		sb.Having(sqlbuilder.Escape(hv))
+	}
+
+	sql, args := sb.BuildWithFlavor(sqlbuilder.ClickHouse)
+	return &qbtypes.Statement{Query: sql, Args: args}, nil
 }
 
 // buildTraceListQuery wires the CTE pipeline (start/end are nanoseconds):

pkg/statementbuilder/tracesstatementbuilder/statement_builder.go

@@ -32,6 +32,9 @@ type traceQueryStatementBuilder struct {
 	resourceFilterResolver         *resourcefilter.ResourceFingerprintResolver[qbtypes.TraceAggregation]
 	aggExprRewriter                qbtypes.AggExprRewriter
 	skipResourceFingerprintEnabled bool
+	// traceScope is set only on the per-call copy made by BuildTraceScoped; it
+	// constrains the query to trace ids selected by the __trace_scope CTE.
+	traceScope *qbtypes.Statement
 }
 
 var _ qbtypes.StatementBuilder[qbtypes.TraceAggregation] = (*traceQueryStatementBuilder)(nil)
@@ -95,6 +98,38 @@ func NewTraceQueryStatementBuilder(
 	}
 }
 
+// BuildTraceScoped is Build additionally constrained to spans whose trace_id is
+// selected by traceScope. The receiver is copied so the shared builder stays stateless.
+func (b *traceQueryStatementBuilder) BuildTraceScoped(
+	ctx context.Context,
+	orgID valuer.UUID,
+	start uint64,
+	end uint64,
+	requestType qbtypes.RequestType,
+	query qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation],
+	variables map[string]qbtypes.VariableItem,
+	traceScope *qbtypes.Statement,
+) (*qbtypes.Statement, error) {
+	// The scope is wired into the list query only; reject other request types rather
+	// than silently dropping the constraint.
+	if requestType != qbtypes.RequestTypeRaw {
+		return nil, errors.NewInternalf(errors.CodeInternal, "trace-scoped build supports only the raw request type, got %s", requestType.StringValue())
+	}
+	scoped := *b
+	scoped.traceScope = traceScope
+	return scoped.Build(ctx, orgID, start, end, requestType, query, variables)
+}
+
+// attachTraceScope adds the trace-scope condition to sb and returns the CTE fragment
+// + args to prepend; both empty when no scope is set.
+func (b *traceQueryStatementBuilder) attachTraceScope(sb *sqlbuilder.SelectBuilder) (string, []any) {
+	if b.traceScope == nil {
+		return "", nil
+	}
+	sb.Where("trace_id GLOBAL IN (SELECT trace_id FROM __trace_scope)")
+	return fmt.Sprintf("__trace_scope AS (%s)", b.traceScope.Query), b.traceScope.Args
+}
+
 // Build builds a SQL query for traces based on the given parameters.
 func (b *traceQueryStatementBuilder) Build(
 	ctx context.Context,
@@ -318,6 +353,11 @@ func (b *traceQueryStatementBuilder) buildListQuery(
 		cteArgs = append(cteArgs, args)
 	}
 
+	if scopeFrag, scopeArgs := b.attachTraceScope(sb); scopeFrag != "" {
+		cteFragments = append(cteFragments, scopeFrag)
+		cteArgs = append(cteArgs, scopeArgs)
+	}
+
 	for i, field := range query.SelectFields {
 		expr, err := b.fm.ColumnExpressionFor(ctx, orgID, start, end, &field, telemetrytypes.FieldDataTypeUnspecified, keys)
 		if err != nil {

tests/integration/tests/querierai/01_ai_traces.py

@@ -214,6 +214,41 @@ def test_ai_span_list_excludes_non_gen_ai_spans(
     assert "POST /api/chat" not in names  # root span excluded
 
 
+def test_ai_span_list_trace_level_filter(
+    signoz: types.SigNoz,
+    create_user_admin: None,  # pylint: disable=unused-argument
+    get_token: Callable[[str, str], str],
+    insert_traces: Callable[[list[Traces]], None],
+) -> None:
+    """Span list (raw) with a trace-level condition returns only the gen_ai spans of
+    traces whose window-clipped aggregates qualify: of two traces with out-tokens
+    100/300, `trace.output_tokens > 100` keeps only the large one's LLM span."""
+    now = datetime.now(tz=UTC).replace(second=0, microsecond=0)
+    service = "ai-it-spanlist-tracefilter"
+    small = ai_trace(now=now, service=service, user="a", in_tokens=10, out_tokens=100)
+    large = ai_trace(now=now, service=service, user="b", in_tokens=30, out_tokens=300)
+    insert_traces(small + large)
+
+    token = get_token(USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD)
+    start_ms, end_ms = query_window(now)
+
+    query = BuilderQuery(
+        signal="traces",
+        query_type="builder_ai_query",
+        name="A",
+        filter_expression=f"service.name = '{service}' AND trace.output_tokens > 100",
+        limit=10,
+    )
+    response = make_query_request(signoz, token, start_ms, end_ms, [query.to_dict()], request_type=RequestType.RAW)
+    assert response.status_code == HTTPStatus.OK, response.text
+
+    rows = response.json()["data"]["data"]["results"][0]["rows"]
+    assert len(rows) == 1, f"expected only the large trace's LLM span, got {len(rows)} rows"
+    body = json.dumps(rows)
+    assert large[0].trace_id in body
+    assert small[0].trace_id not in body
+
+
 def test_ai_list_having_or_aggregates(
     signoz: types.SigNoz,
     create_user_admin: None,  # pylint: disable=unused-argument

tests/integration/tests/querierai/conftest.py

@@ -0,0 +1,38 @@
+import pytest
+from testcontainers.core.container import Network
+
+from fixtures import types
+from fixtures.signoz import create_signoz
+
+
+@pytest.fixture(name="signoz", scope="package")
+def signoz_ai_observability(
+    network: Network,
+    migrator: types.Operation,  # pylint: disable=unused-argument
+    zeus: types.TestContainerDocker,
+    gateway: types.TestContainerDocker,
+    sqlstore: types.TestContainerSQL,
+    clickhouse: types.TestContainerClickhouse,
+    request: pytest.FixtureRequest,
+    pytestconfig: pytest.Config,
+) -> types.SigNoz:
+    """
+    Package-scoped SigNoz instance with AI observability enabled. builder_ai_query
+    relies on the metadata store surfacing the static gen_ai key definitions
+    (enrichWithGenAIKeys), which is gated on this flag — without it the gate keys
+    (gen_ai.tool.name, gen_ai.agent.name, ...) only resolve once a span carrying
+    them has been ingested.
+    """
+    return create_signoz(
+        network=network,
+        zeus=zeus,
+        gateway=gateway,
+        sqlstore=sqlstore,
+        clickhouse=clickhouse,
+        request=request,
+        pytestconfig=pytestconfig,
+        cache_key="signoz-ai-observability",
+        env_overrides={
+            "SIGNOZ_FLAGGER_CONFIG_BOOLEAN_ENABLE__AI__OBSERVABILITY": True,
+        },
+    )