PR #62 follow-up 14: signal get_trace_activity truncation at 500 rows

docs/tool-api.md

@@ -148,7 +148,8 @@ top-level set operation such as `UNION` and to a parenthesized outer query.
 Projections inside CTE definitions, derived tables and subqueries are not
 checked. Casts to `INT64`, `UINT64` or `FLOAT64` anywhere in the statement are
 rejected with a hint to use `BIGINT` or `DOUBLE`. The statement is wrapped with
-an outer `LIMIT 500`. `sample_events`
+an outer `LIMIT 500`, except in `get_trace_activity`, whose generated SQL carries
+its own limit (see below). `sample_events`
 returns at most 10 columns, choosing common observability fields such as
 `p_timestamp`, `service.name`, `severity_text`, `body` and `trace_id` before
 other fields.
@@ -243,14 +244,17 @@ dataset's time partition, or `p_timestamp` if none is configured. The dataset
 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. Rows are ordered by bucket. The result contains `query`,
-`bucketInterval`, `semantics` and up to 500 `rows`.
-
-`query` is the aggregation SQL before the tool wraps it with `LIMIT 500`. It
-has no time predicate: the time range comes from the request's `startTime` and
-`endTime`, which the query API rounds down to the minute. To reproduce the
-result, for example in a dashboard tile, run `query` through the query API with
-the same `startTime` and `endTime`.
+not double counted. Rows are ordered by bucket, oldest first. The result
+contains `query`, `bucketInterval`, `semantics`, up to 500 `rows` and
+`truncated`.
+
+`query` is the exact SQL the tool runs; it ends in `LIMIT 501` so the tool can
+detect overflow. When more than 500 rows match, the newest buckets are dropped
+and `truncated` is `true`; reduce `bucketCount`, `groupBy` or the time range to
+get the full series. `query` has no time predicate: the time range comes from
+the request's `startTime` and `endTime`, which the query API rounds down to the
+minute. To reproduce the result, for example in a dashboard tile, run `query`
+through the query API with the same `startTime` and `endTime`.
 
 ## Argument schema validation
 

src/tool_executor.rs

@@ -1400,15 +1400,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<'_>,
@@ -1532,9 +1549,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,
@@ -1547,7 +1567,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",
@@ -1558,6 +1579,7 @@ async fn execute_trace_activity_tool(
             "trace_count": "COUNT(DISTINCT span_trace_id)"
         },
         "rows": rows,
+        "truncated": truncated,
     }))
 }
 
@@ -2305,21 +2327,8 @@ 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 mut writer = arrow_json::ArrayWriter::new(Vec::new());
-        writer
-            .write_batches(&batches.iter().collect::<Vec<_>>())
-            .unwrap();
-        writer.finish().unwrap();
-        let rows: Vec<Value> = serde_json::from_slice(&writer.into_inner()).unwrap();
+        let query = trace_activity_query("traces", "p_timestamp", &interval, group_by, clauses);
+        let rows = execute_trace_activity_sql(&context, &query).await;
         let times = rows
             .iter()
             .map(|row| row["time"].as_str().unwrap())
@@ -2331,6 +2340,69 @@ mod tests {
         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<Value> {
+        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::<Vec<_>>())
+            .unwrap();
+        writer.finish().unwrap();
+        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]
     async fn trace_activity_query_executes_with_distinct_bucketed_counts() {
         let rows = run_trace_activity_fixture(&["service.name".to_owned()], &[], "BIGINT").await;