From 2beb4bc97adb9211e693cdc4e5be63b7d7dfb0ab Mon Sep 17 00:00:00 2001 From: VibeSWE <338779468+VibeSWE@users.noreply.github.com> Date: Sat, 24 Oct 2026 14:30:53 -0500 Subject: [PATCH] Signal truncation in get_trace_activity instead of dropping buckets silently get_trace_activity ordered rows by bucket time and capped them at 500, both in the generated SQL and in the generic tool wrapper. With high groupBy cardinality (e.g. 20 services x 60 buckets) the newest buckets were cut off with no indication, which reads as traffic dropping to zero. The generated query now fetches MAX_TOOL_ROWS + 1 rows and runs as-is after validation (it already carries its own LIMIT), so the tool can return at most 500 rows plus a `truncated` flag. The returned `query` is now exactly the SQL that ran. Co-Authored-By: Claude Opus 5.5 --- docs/tool-api.md | 7 ++- src/tool_executor.rs | 102 ++++++++++++++++++++++++++++++++++++------- 2 files changed, 93 insertions(+), 16 deletions(-) diff --git a/docs/tool-api.md b/docs/tool-api.md index fc9cca74f..8303ca702 100644 --- a/docs/tool-api.md +++ b/docs/tool-api.md @@ -212,9 +212,14 @@ must have traces telemetry and contain `span_span_id` and `span_trace_id`. Unknown fields are rejected. `span_count` is `COUNT(DISTINCT span_span_id)` and `trace_count` is `COUNT(DISTINCT span_trace_id)`, so duplicate span rows are not double counted. The result contains the exact executed SQL as `query`, -along with `bucketInterval`, `semantics` and up to 500 `rows`. Clients can +along with `bucketInterval`, `semantics`, up to 500 `rows` and `truncated`. Clients can rerun that SQL unchanged, for example in a dashboard tile. +Rows are ordered by bucket time, oldest first. The SQL fetches one row past the +limit (`LIMIT 501`) to detect overflow. When more than 500 rows match, the +newest are dropped and `truncated` is `true`; reduce `bucketCount`, `groupBy` +or the time range to get the full series. + ## Argument schema validation `inputSchema` is enforced server-side for `type` (a single type or a list of types), `required`, diff --git a/src/tool_executor.rs b/src/tool_executor.rs index f2c477041..e55efc792 100644 --- a/src/tool_executor.rs +++ b/src/tool_executor.rs @@ -1378,15 +1378,32 @@ fn trace_activity_query( .first() .map(|_| format!(" WHERE {}", clauses.join(" AND "))) .unwrap_or_default(); + // One row past the tool limit lets the caller detect truncation. format!( - "SELECT {} FROM {}{} GROUP BY {} ORDER BY 1 LIMIT 500", + "SELECT {} FROM {}{} GROUP BY {} ORDER BY 1 LIMIT {}", select.join(", "), quote_identifier(dataset), where_clause, groups, + MAX_TOOL_ROWS + 1, ) } +/// Caps trace activity rows at the tool limit and reports whether any were dropped. +fn truncate_trace_activity_rows(rows: &mut Value) -> bool { + let rows = match rows { + Value::Array(rows) => rows, + Value::Object(object) => match object.get_mut("records").and_then(Value::as_array_mut) { + Some(rows) => rows, + None => return false, + }, + _ => return false, + }; + let truncated = rows.len() > MAX_TOOL_ROWS as usize; + rows.truncate(MAX_TOOL_ROWS as usize); + truncated +} + async fn trace_activity_fields( dataset: &str, context: OssToolContext<'_>, @@ -1501,9 +1518,12 @@ async fn execute_trace_activity_tool( let bucket_count = bounded_integer(arguments, "bucketCount", 60, 1, 100); let interval = time_bucket_interval(start_time, end_time, bucket_count); let query = trace_activity_query(dataset, &time_field, &interval, &group_by, &clauses); + // The generated query carries its own LIMIT; wrapping it in the generic + // 500-row bound would hide the extra row used to detect truncation. + validate_tool_sql(&query)?; let response = query::query_internal( query::Query { - query: bounded_select(&query)?, + query: query.clone(), start_time: start_time.to_owned(), end_time: end_time.to_owned(), send_null: true, @@ -1516,7 +1536,8 @@ async fn execute_trace_activity_tool( ) .await .map_err(|error| error.to_string())?; - let rows = response_value(response).await?; + let mut rows = response_value(response).await?; + let truncated = truncate_trace_activity_rows(&mut rows); Ok(json!({ "dataset": dataset, "queryType": "sql", @@ -1527,6 +1548,7 @@ async fn execute_trace_activity_tool( "trace_count": "COUNT(DISTINCT span_trace_id)" }, "rows": rows, + "truncated": truncated, })) } @@ -2117,23 +2139,73 @@ mod tests { let interval = super::time_bucket_interval("2026-10-22T00:00:00Z", "2026-10-22T00:02:00Z", 2); assert_eq!(interval, "60 seconds"); - let query = bounded_tool_sql(&trace_activity_query( - "traces", - "p_timestamp", - &interval, - group_by, - clauses, - )) - .unwrap(); - let batches = context.sql(&query).await.unwrap().collect().await.unwrap(); + let query = trace_activity_query("traces", "p_timestamp", &interval, group_by, clauses); + let mut rows = execute_trace_activity_sql(&context, &query).await; + rows.sort_by_key(|row| row.to_string()); + rows + } + + /// Runs trace activity SQL the way the tool does: validated, without the + /// generic 500-row wrapper. + async fn execute_trace_activity_sql( + context: &datafusion::prelude::SessionContext, + query: &str, + ) -> Vec { + super::validate_tool_sql(query).unwrap(); + let batches = context.sql(query).await.unwrap().collect().await.unwrap(); let mut writer = arrow_json::ArrayWriter::new(Vec::new()); writer .write_batches(&batches.iter().collect::>()) .unwrap(); writer.finish().unwrap(); - let mut rows: Vec = serde_json::from_slice(&writer.into_inner()).unwrap(); - rows.sort_by_key(|row| row.to_string()); - rows + serde_json::from_slice(&writer.into_inner()).unwrap() + } + + #[actix_web::test] + async fn trace_activity_signals_truncation_past_the_row_limit() { + // 20 services x 60 one-minute buckets = 1200 groups x buckets. + let context = datafusion::prelude::SessionContext::new(); + let fixture = context + .sql( + r#"SELECT CAST(TO_TIMESTAMP_SECONDS(1792627200 + (value / 20) * 60) AS TIMESTAMP) + AS p_timestamp, + CONCAT('t', value) AS span_trace_id, CONCAT('s', value) AS span_span_id, + CONCAT('svc-', value % 20) AS "service.name" + FROM range(1200)"#, + ) + .await + .unwrap(); + context + .register_table("traces", fixture.into_view()) + .unwrap(); + let interval = + super::time_bucket_interval("2026-10-22T00:00:00Z", "2026-10-22T01:00:00Z", 60); + assert_eq!(interval, "60 seconds"); + let query = trace_activity_query( + "traces", + "p_timestamp", + &interval, + &["service.name".to_owned()], + &[], + ); + let mut rows = Value::Array(execute_trace_activity_sql(&context, &query).await); + assert!(super::truncate_trace_activity_rows(&mut rows)); + let rows = rows.as_array().unwrap(); + assert_eq!(rows.len(), 500); + // Ascending order keeps the oldest 25 buckets; the newest are the ones dropped. + assert_eq!(rows[499]["time"], "2026-10-22T00:24:00"); + + // A result that fits is returned whole and not flagged. + let query = trace_activity_query("traces", "p_timestamp", &interval, &[], &[]); + let mut rows = Value::Array(execute_trace_activity_sql(&context, &query).await); + assert!(!super::truncate_trace_activity_rows(&mut rows)); + assert_eq!(rows.as_array().unwrap().len(), 60); + + let mut exact = Value::Array(vec![json!({}); 500]); + assert!(!super::truncate_trace_activity_rows(&mut exact)); + let mut records = json!({ "records": vec![json!({}); 501] }); + assert!(super::truncate_trace_activity_rows(&mut records)); + assert_eq!(records["records"].as_array().unwrap().len(), 500); } #[actix_web::test]