From 2f8718daa49c3e88c36e9337989faf1aecc7b4c1 Mon Sep 17 00:00:00 2001 From: Nikhil Sinha <131262146+nikhilsinhaparseable@users.noreply.github.com> Date: Thu, 22 Oct 2026 14:04:25 +0530 Subject: [PATCH 1/5] update: tool catalog (#1831) Port of upstream parseablehq/parseable 32979cb4. add get_trace_activity to tool catalog update executor to have tool implementations and reuse in enterprise Fork conflict resolutions: - tool_catalog: use upstream's validate_one_of helper (same exactly-one semantics and message as the fork's inline check) and keep the fork's `const` validation; tool counts become 54 total / 41 read. - tool_executor: adopt upstream's parser-based validate_tool_sql and bounded_tool_sql; keep fork time-expression imports, the Result handling for list_dashboards_internal, and the fork's call_tool_registry (dashboard tool routing, suspension and Query gates); keep calling execute_oss_tool_inner there because the spec is already validated. - tests: merge both suites; update the fork bounded-SQL test to avoid a wildcard CTE that upstream's validator now rejects. Also add tests for get_trace_activity (executed against a DataFusion fixture), its argument validation, resolve_dataset/get_log_context schema changes, timezone day boundaries and sample_events range pairing, and document the new and changed tools in docs/tool-api.md. (cherry picked from commit 32979cb4cf904296a9d5d88f400d00df358cef8b) Co-Authored-By: Claude Opus 5.5 --- Cargo.lock | 1 + Cargo.toml | 1 + docs/tool-api.md | 92 ++ src/handlers/http/query_context.rs | 26 +- src/tool_catalog.rs | 251 ++++- src/tool_executor.rs | 1473 ++++++++++++++++++++++++---- 6 files changed, 1597 insertions(+), 247 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 386aced06..0726e1b08 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4730,6 +4730,7 @@ dependencies = [ "bytes", "chrono", "chrono-humanize", + "chrono-tz", "clap", "clokwerk", "cookie 0.18.1", diff --git a/Cargo.toml b/Cargo.toml index d2cc8d8d3..f8c6459f9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -139,6 +139,7 @@ tracing-actix-web = "0.7" # Time and Date chrono = "0.4" +chrono-tz = "0.10" chrono-humanize = "0.2" humantime = "2.1.0" humantime-serde = "1.1" diff --git a/docs/tool-api.md b/docs/tool-api.md index 73e0241e3..b57ed774d 100644 --- a/docs/tool-api.md +++ b/docs/tool-api.md @@ -130,3 +130,95 @@ and last-seen value, and coalesces supported aliases per row. Its `total` is the number of returned groups, not an unbounded count of all matching groups. The native Errors page additionally separates issues by message, so its grouping and pagination are not identical to this summary tool. + +## SQL tool limits + +`query_sql`, `explain_query` and other tools that run SQL internally parse the +statement with DataFusion's SQL parser and accept exactly one `SELECT` or `WITH` +query. A trailing semicolon and trailing comments are allowed. Every `SELECT` +projection, including those inside nested queries and set operations, must +name at most 10 explicit columns. `*` and `table.*` are rejected. Casts to +`INT64`, `UINT64` or `FLOAT64` are rejected with a hint to use `BIGINT` or +`DOUBLE`. The statement is wrapped with an outer `LIMIT 500`. `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. + +## Dataset and time resolution + +`list_datasets`, `resolve_dataset` and `discover_datasets` omit internal +datasets. + +`resolve_dataset` accepts an optional `telemetryType` (`logs`, `metrics` or +`traces`) and `limit` (1–20, default 5). It resolves a `selectedDataset` or +`contextDataset` exactly, or resolves the only visible dataset. Otherwise it +ranks visible datasets by name match against `query`. It resolves only a +unique high-confidence name match. In every other case it returns +`status: "ambiguous"` with up to `limit` scored `candidates`, or `"not_found"`. +It never falls back to an arbitrary dataset. A selected dataset whose telemetry +type does not match `telemetryType` returns `not_found`. + +`discover_datasets` counts events in the requested range for each visible +dataset of the requested telemetry type that the caller can query. It returns +active datasets ordered by event count as `{ name, eventsInRange }`, together +with `selectedDatasets`, `activeDatasetCount`, `truncated` and up to 20 +`unavailableDatasets`. If no dataset is active and at least one could not be +counted, the result has `status: "error"` and no datasets. + +`resolve_time_range` accepts any IANA `timezone` (default `UTC`), and uses it +for `today` and `yesterday` day boundaries. It also accepts `last `, +`past `, `now-`, ` to ` and RFC3339 start +times. An unknown timezone is an error. + +`sample_events` accepts an explicit `startTime` and `endTime`, which must be +provided together. They take precedence over `minutes`. + +## Log context + +`get_log_context` uses the same implementation as `POST /api/v1/query/context`. +Arguments must include exactly one of `log`, `body` or `message`, holding the +anchor row's exact value. The schema enforces this with `oneOf`, so calls with +none or more than one are rejected with `400`. If no context bounds are given, +the window defaults to `contextWindow: "5 minutes"`. Returned records keep at +most 10 fields per row, preferring `p_timestamp`, service, severity, message +and trace identifiers. The response also includes the anchor index, scope and +the executed queries. + +## Trace activity + +`get_trace_activity` returns a bucketed span and trace activity series for one +traces dataset: + +```json +{ + "name": "get_trace_activity", + "arguments": { + "dataset": "otel-traces", + "startTime": "2026-10-22T00:00:00Z", + "endTime": "2026-10-22T01:00:00Z", + "groupBy": ["service.name"], + "filters": { "span_status_code": 2 }, + "bucketCount": 60 + } +} +``` + +`dataset`, `startTime` and `endTime` are required. `groupBy` takes up to 5 +exact field names. `filters` maps exact field names to JSON scalars: strings, +numbers and booleans test equality, and `null` tests `IS NULL`. `bucketCount` +is 1–100, default 60. It approximates the number of `DATE_BIN` buckets over the +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. The result contains the exact executed SQL as `query`, +along with `bucketInterval`, `semantics` and up to 500 `rows`. Clients can +rerun that SQL unchanged, for example in a dashboard tile. + +## Argument schema validation + +`inputSchema` is enforced server-side for `type`, `required`, +`additionalProperties`, `enum`, `const`, `minLength`/`maxLength`, +`minimum`/`maximum`, `minItems`/`maxItems`, `items`, `anyOf` and `oneOf`. +`oneOf` passes only when exactly one variant matches. Combinators are checked +in addition to sibling keywords; they do not replace them. diff --git a/src/handlers/http/query_context.rs b/src/handlers/http/query_context.rs index 3b4229c29..64df0039b 100644 --- a/src/handlers/http/query_context.rs +++ b/src/handlers/http/query_context.rs @@ -138,7 +138,6 @@ pub async fn query_context( let permissions = Users.get_permissions(&creds); create_streams_for_distributed(vec![context_request.dataset.clone()], &tenant_id).await?; - let authorized_datasets = vec![context_request.dataset.clone()]; user_auth_for_datasets( &permissions, std::slice::from_ref(&context_request.dataset), @@ -146,6 +145,19 @@ pub async fn query_context( ) .await?; + query_context_for_authorized_dataset(context_request, &tenant_id) + .await + .map(web::Json) +} + +/// Fetches log context after the caller has authorized the dataset and made it +/// available on the local query node. +pub(crate) async fn query_context_for_authorized_dataset( + context_request: LogContextRequest, + tenant_id: &Option, +) -> Result { + let span = Span::current(); + let authorized_datasets = vec![context_request.dataset.clone()]; let page_size = normalize_log_context_page_size(context_request.page_size)?; span.record("page_size", page_size); let anchor_timestamp = parse_log_context_timestamp(&context_request.p_timestamp)?; @@ -164,7 +176,7 @@ pub async fn query_context( "query context request normalized" ); - let dataset = PARSEABLE.get_stream(&context_request.dataset, &tenant_id)?; + let dataset = PARSEABLE.get_stream(&context_request.dataset, tenant_id)?; let schema = dataset.get_schema(); validate_log_context_schema(schema.as_ref(), DEFAULT_TIMESTAMP_KEY)?; let match_fields = normalize_log_context_match_fields( @@ -191,7 +203,7 @@ pub async fn query_context( &authorized_datasets, &context_start_time_str, &context_end_time_str, - &tenant_id, + tenant_id, ) .await?; debug!( @@ -242,8 +254,8 @@ pub async fn query_context( ); let (newer_records, older_records) = tokio::try_join!( - execute_log_context_rows(&newer_payload, &authorized_datasets, &tenant_id), - execute_log_context_rows(&older_payload, &authorized_datasets, &tenant_id), + execute_log_context_rows(&newer_payload, &authorized_datasets, tenant_id), + execute_log_context_rows(&older_payload, &authorized_datasets, tenant_id), )?; let (records, anchor_index) = build_log_context_records_window(newer_records, older_records, page_size)?; @@ -274,7 +286,7 @@ pub async fn query_context( tenant_id.as_deref().unwrap_or(DEFAULT_TENANT), ); - Ok(web::Json(LogContextResponse { + Ok(LogContextResponse { scope, context_start_time: context_start_time_str, context_end_time: context_end_time_str, @@ -284,7 +296,7 @@ pub async fn query_context( anchored_duplicate: LOG_CONTEXT_ANCHORED_DUPLICATE, records, queries, - })) + }) } fn normalize_log_context_page_size(page_size: Option) -> Result { diff --git a/src/tool_catalog.rs b/src/tool_catalog.rs index 1ec3625aa..8edc2f833 100644 --- a/src/tool_catalog.rs +++ b/src/tool_catalog.rs @@ -106,6 +106,7 @@ impl ToolSpec { | "discover_datasets" | "get_traces" | "get_trace" + | "get_trace_activity" | "list_log_errors" | "list_trace_errors" => ToolEffect::ExpensiveRead, "create_alert" @@ -138,14 +139,8 @@ fn validate_schema_value(schema: &Value, value: &Value, path: &str) -> Result<() if let Some(result) = validate_any_of(schema, value, path) { result?; } - if let Some(variants) = schema.get("oneOf").and_then(Value::as_array) { - let matches = variants - .iter() - .filter(|variant| validate_schema_value(variant, value, path).is_ok()) - .count(); - if matches != 1 { - return Err(format!("{path} must match exactly one allowed schema")); - } + if let Some(result) = validate_one_of(schema, value, path) { + result?; } if let Some(expected) = schema.get("const") && expected != value @@ -161,6 +156,21 @@ fn validate_schema_value(schema: &Value, value: &Value, path: &str) -> Result<() validate_number(schema, value, path) } +fn validate_one_of(schema: &Value, value: &Value, path: &str) -> Option> { + schema + .get("oneOf") + .and_then(Value::as_array) + .map(|variants| { + let matches = variants + .iter() + .filter(|variant| validate_schema_value(variant, value, path).is_ok()) + .count(); + (matches == 1) + .then_some(()) + .ok_or_else(|| format!("{path} must match exactly one allowed schema")) + }) +} + fn validate_any_of(schema: &Value, value: &Value, path: &str) -> Option> { schema .get("anyOf") @@ -357,7 +367,9 @@ pub fn oss_read_tool_specs() -> Vec { json!({ "query": string_property("User request or dataset search text."), "selectedDataset": { "type": "string" }, - "contextDataset": { "type": "string" } + "contextDataset": { "type": "string" }, + "telemetryType": { "type": "string", "enum": ["logs", "metrics", "traces"] }, + "limit": { "type": "integer", "minimum": 1, "maximum": 20, "default": 5 } }), &["query"], ), @@ -407,6 +419,8 @@ pub fn oss_read_tool_specs() -> Vec { "dataset": string_property("Dataset name."), "limit": { "type": "integer", "minimum": 1, "maximum": 100, "default": 10 }, "minutes": { "type": "integer", "minimum": 1, "maximum": 1440, "default": 60 }, + "startTime": string_property("Optional explicit RFC3339 range start; provide with endTime."), + "endTime": string_property("Optional explicit RFC3339 range end; provide with startTime."), }), &["dataset"], ), @@ -415,21 +429,29 @@ pub fn oss_read_tool_specs() -> Vec { "get_log_context", "Get surrounding log context", "Fetch records around an exact anchor log row.", - object_schema( - json!({ - "dataset": string_property("Dataset containing the anchor row."), - "pTimestamp": string_property("Exact p_timestamp from the anchor row."), - "log": { "type": "string" }, - "body": { "type": "string" }, - "message": { "type": "string" }, - "contextWindow": { "type": "string" }, - "contextStartTime": { "type": "string" }, - "contextEndTime": { "type": "string" }, - "conditions": { "type": "object", "additionalProperties": true }, - "pageSize": { "type": "integer", "minimum": 1, "maximum": 500, "default": 100 }, - }), - &["dataset", "pTimestamp"], - ), + { + let mut schema = object_schema( + json!({ + "dataset": string_property("Dataset containing the anchor row."), + "pTimestamp": string_property("Exact p_timestamp from the anchor row."), + "log": string_property("Exact log value from the anchor row. Provide exactly one of log, body, or message."), + "body": string_property("Exact body value from the anchor row. Provide exactly one of log, body, or message."), + "message": string_property("Exact message value from the anchor row. Provide exactly one of log, body, or message."), + "contextWindow": { "type": "string" }, + "contextStartTime": { "type": "string" }, + "contextEndTime": { "type": "string" }, + "conditions": { "type": "object", "additionalProperties": true }, + "pageSize": { "type": "integer", "minimum": 1, "maximum": 500, "default": 100 }, + }), + &["dataset", "pTimestamp"], + ); + schema["oneOf"] = json!([ + { "required": ["log"] }, + { "required": ["body"] }, + { "required": ["message"] } + ]); + schema + }, ), ToolSpec::new( "list_log_errors", @@ -531,6 +553,22 @@ pub fn oss_read_tool_specs() -> Vec { &["traceId", "dataset", "startTime", "endTime"], ), ), + ToolSpec::new( + "get_trace_activity", + "Get trace activity", + "Return a bounded trace-activity time series using Parseable's canonical semantics: span_count is COUNT(DISTINCT span_span_id) and trace_count is COUNT(DISTINCT span_trace_id). Use this for span or trace volume, trends, charts, and service breakdowns instead of deriving counts from raw trace rows. Returns the exact executed SQL with its rows so clients can reuse it unchanged.", + object_schema( + json!({ + "dataset": string_property("One traces dataset."), + "startTime": string_property("RFC3339 range start."), + "endTime": string_property("RFC3339 range end."), + "groupBy": { "type": "array", "maxItems": 5, "items": { "type": "string" }, "description": "Optional exact trace fields, such as service.name, used to split the activity series." }, + "filters": { "type": "object", "description": "Optional exact trace-field filters.", "additionalProperties": true }, + "bucketCount": { "type": "integer", "minimum": 1, "maximum": 100, "default": 60, "description": "Approximate number of time buckets." } + }), + &["dataset", "startTime", "endTime"], + ), + ), ToolSpec::new( "ping", "Ping Parseable server", @@ -848,7 +886,7 @@ mod tests { let tools = oss_tool_specs(); let names = tools.iter().map(|tool| tool.name).collect::>(); - assert_eq!(tools.len(), 53); + assert_eq!(tools.len(), 54); assert_eq!(tools.len(), names.len()); } @@ -887,7 +925,7 @@ mod tests { let read_tools = oss_read_tool_specs(); let action_tools = oss_action_tool_specs(); - assert_eq!(read_tools.len(), 40); + assert_eq!(read_tools.len(), 41); assert_eq!(action_tools.len(), 13); assert!(read_tools.iter().all(|tool| matches!( tool.effect(), @@ -971,6 +1009,121 @@ mod tests { ); } + #[actix_web::test] + async fn oss_catalog_exposes_trace_activity() { + let tool = oss_read_tool_specs() + .into_iter() + .find(|tool| tool.name == "get_trace_activity") + .expect("get_trace_activity must be available in OSS"); + + assert_eq!(tool.effect(), ToolEffect::ExpensiveRead); + assert_eq!(tool.input_schema["properties"]["dataset"]["type"], "string"); + assert_eq!(tool.input_schema["properties"]["groupBy"]["maxItems"], 5); + + let range = serde_json::json!({ + "dataset": "traces", + "startTime": "2026-10-22T00:00:00Z", + "endTime": "2026-10-22T01:00:00Z" + }); + assert!(tool.validate_arguments(&range).is_ok()); + let mut grouped = range.clone(); + grouped["groupBy"] = serde_json::json!(["service.name"]); + grouped["filters"] = serde_json::json!({ "span_status_code": 2 }); + grouped["bucketCount"] = serde_json::json!(100); + assert!(tool.validate_arguments(&grouped).is_ok()); + for (field, value) in [ + ("groupBy", serde_json::json!(["a", "b", "c", "d", "e", "f"])), + ("groupBy", serde_json::json!([1])), + ("bucketCount", serde_json::json!(0)), + ("bucketCount", serde_json::json!(101)), + ("unexpected", serde_json::json!(true)), + ] { + let mut invalid = range.clone(); + invalid[field] = value; + assert!(tool.validate_arguments(&invalid).is_err(), "{field}"); + } + for missing in ["dataset", "startTime", "endTime"] { + let mut invalid = range.clone(); + invalid.as_object_mut().unwrap().remove(missing); + assert!(tool.validate_arguments(&invalid).is_err(), "{missing}"); + } + } + + #[actix_web::test] + async fn registry_lists_trace_activity_as_read_only() { + let app = + test::init_service(App::new().route("/tools/list", web::get().to(list_tool_registry))) + .await; + let response: Value = test::call_and_read_body_json( + &app, + test::TestRequest::get().uri("/tools/list").to_request(), + ) + .await; + let tool = response["tools"] + .as_array() + .unwrap() + .iter() + .find(|tool| tool["name"] == "get_trace_activity") + .expect("get_trace_activity must be externally listed"); + + assert_eq!(tool["annotations"]["readOnlyHint"], true); + assert_eq!(tool["annotations"]["destructiveHint"], false); + } + + #[actix_web::test] + async fn changed_read_tool_schemas_reject_out_of_contract_arguments() { + let find = |name: &str| { + oss_read_tool_specs() + .into_iter() + .find(|tool| tool.name == name) + .unwrap() + }; + let resolve = find("resolve_dataset"); + for invalid in [ + serde_json::json!({ "query": "x", "telemetryType": "events" }), + serde_json::json!({ "query": "x", "limit": 0 }), + serde_json::json!({ "query": "x", "limit": 21 }), + ] { + assert!(resolve.validate_arguments(&invalid).is_err(), "{invalid}"); + } + + let context = find("get_log_context"); + for anchor in ["log", "body", "message"] { + assert!( + context + .validate_arguments(&serde_json::json!({ + "dataset": "logs", + "pTimestamp": "2026-10-22T00:00:00Z", + anchor: "failed" + })) + .is_ok(), + "{anchor}" + ); + } + assert!( + context + .validate_arguments(&serde_json::json!({ + "dataset": "logs", + "pTimestamp": "2026-10-22T00:00:00Z", + "log": "a", + "body": "b", + "message": "c" + })) + .is_err() + ); + + let sample = find("sample_events"); + assert!( + sample + .validate_arguments(&serde_json::json!({ + "dataset": "logs", + "startTime": "2026-10-22T00:00:00Z", + "endTime": "2026-10-22T01:00:00Z" + })) + .is_ok() + ); + } + #[actix_web::test] async fn tool_call_defaults_to_empty_arguments() { let call: ToolCallRequest = serde_json::from_value(serde_json::json!({ @@ -1028,6 +1181,52 @@ mod tests { .is_err() ); } + + let resolve_tool = oss_tool_specs() + .into_iter() + .find(|tool| tool.name == "resolve_dataset") + .unwrap(); + assert!( + resolve_tool + .validate_arguments(&serde_json::json!({ + "query": "metrics", + "telemetryType": "metrics", + "limit": 10 + })) + .is_ok() + ); + + let context_tool = oss_tool_specs() + .into_iter() + .find(|tool| tool.name == "get_log_context") + .unwrap(); + assert!( + context_tool + .validate_arguments(&serde_json::json!({ + "dataset": "logs", + "pTimestamp": "2026-10-22T00:00:00Z" + })) + .is_err() + ); + assert!( + context_tool + .validate_arguments(&serde_json::json!({ + "dataset": "logs", + "pTimestamp": "2026-10-22T00:00:00Z", + "message": "failed" + })) + .is_ok() + ); + assert!( + context_tool + .validate_arguments(&serde_json::json!({ + "dataset": "logs", + "pTimestamp": "2026-10-22T00:00:00Z", + "log": "failed", + "message": "failed" + })) + .is_err() + ); } #[actix_web::test] diff --git a/src/tool_executor.rs b/src/tool_executor.rs index 772acac3d..f2c477041 100644 --- a/src/tool_executor.rs +++ b/src/tool_executor.rs @@ -1,21 +1,32 @@ -use std::collections::HashMap; +use std::{collections::HashMap, ops::ControlFlow}; use actix_web::{HttpRequest, HttpResponse, body::to_bytes_limited, http::header::HeaderMap, web}; -use chrono::{DateTime, Duration, Utc}; +use chrono::{DateTime, Duration, LocalResult, TimeZone, Utc}; +use chrono_tz::Tz; +use futures::{StreamExt, stream}; use serde_json::{Value, json}; use ulid::Ulid; use crate::{ + datafusion::sql::{ + parser::{DFParser, Statement as DFStatement}, + sqlparser::ast::{ + DataType as SqlDataType, Expr as SqlExpr, SelectItem, SetExpr, + Statement as SqlStatement, visit_expressions, + }, + }, handlers::http::{ alerts, cluster, health_check, logstream, modal::query::querier_logstream, - query, rbac, role, targets, + query, + query_context::{self, LogContextRequest}, + rbac, role, targets, users::{dashboards, filters}, }, parseable::PARSEABLE, prism::logstream::get_prism_logstream_info, rbac::{Response, Users, map::SessionKey, role::Action}, - tool_catalog::{ToolCallRequest, ToolCallResult}, + tool_catalog::{ToolCallRequest, ToolCallResult, oss_tool_specs}, utils::{ actix::extract_session_key_from_req, get_tenant_id_from_request, @@ -47,26 +58,120 @@ fn quote_literal(value: &str) -> String { format!("'{}'", value.replace('\'', "''")) } -fn bounded_select(query: &str) -> Result { - let query = query.trim(); - let query = query.strip_suffix(';').unwrap_or(query); - if query.contains(';') { +const MAX_TOOL_ROWS: u64 = 500; +const MAX_TOOL_COLUMNS: usize = 10; +const DATASET_DISCOVERY_CONCURRENCY: usize = 8; + +pub fn validate_tool_sql(query: &str) -> Result<(), String> { + let statements = + DFParser::parse_sql(query.trim()).map_err(|error| format!("invalid SQL: {error}"))?; + if statements.len() != 1 { return Err("exactly one SQL statement is required".to_owned()); } - let first_word = query.split_whitespace().next().unwrap_or(""); - if !(first_word.eq_ignore_ascii_case("select") || first_word.eq_ignore_ascii_case("with")) { + let Some(DFStatement::Statement(statement)) = statements.front() else { return Err("only SQL SELECT queries are allowed".to_owned()); + }; + let SqlStatement::Query(query) = statement.as_ref() else { + return Err("only SQL SELECT queries are allowed".to_owned()); + }; + let unsupported_cast = match visit_expressions(query, |expression| { + let SqlExpr::Cast { data_type, .. } = expression else { + return ControlFlow::Continue(()); + }; + let replacement = match data_type { + SqlDataType::Int64 | SqlDataType::UInt64 => "BIGINT", + SqlDataType::Float64 => "DOUBLE", + _ => return ControlFlow::Continue(()), + }; + ControlFlow::Break((data_type.to_string(), replacement)) + }) { + ControlFlow::Break(unsupported) => Some(unsupported), + ControlFlow::Continue(()) => None, + }; + if let Some((unsupported, replacement)) = unsupported_cast { + return Err(format!( + "unsupported DataFusion SQL cast type `{unsupported}`; use `{replacement}` instead" + )); } + fn validate_projection(body: &SetExpr) -> Result<(), String> { + match body { + SetExpr::Select(select) => { + if select.projection.len() > MAX_TOOL_COLUMNS { + return Err(format!( + "query selects {} columns; the tool limit is {MAX_TOOL_COLUMNS}", + select.projection.len() + )); + } + if select.projection.iter().any(|item| { + matches!( + item, + SelectItem::Wildcard(_) | SelectItem::QualifiedWildcard(_, _) + ) + }) { + return Err(format!( + "wildcard selection is not allowed; select at most {MAX_TOOL_COLUMNS} explicit columns" + )); + } + Ok(()) + } + SetExpr::Query(query) => validate_projection(query.body.as_ref()), + SetExpr::SetOperation { left, right, .. } => { + validate_projection(left)?; + validate_projection(right) + } + _ => Err("unsupported SQL query form".to_owned()), + } + } + validate_projection(query.body.as_ref()) +} + +pub fn bounded_tool_sql(query: &str) -> Result { + validate_tool_sql(query)?; Ok(format!( - "SELECT * FROM (\n{query}\n) AS __parseable_tool_query LIMIT 500" + "SELECT * FROM ({}\n) AS \"__parseable_tool_result\" LIMIT {MAX_TOOL_ROWS}", + query.trim().trim_end_matches(';') )) } +pub fn enforce_tool_tabular_limits(mut value: Value) -> Result { + let rows = match &mut value { + Value::Array(rows) => rows, + Value::Object(object) => object + .get_mut("records") + .and_then(Value::as_array_mut) + .ok_or_else(|| "query response is not tabular".to_owned())?, + _ => return Err("query response is not tabular".to_owned()), + }; + rows.truncate(MAX_TOOL_ROWS as usize); + if let Some(column_count) = rows + .iter() + .filter_map(Value::as_object) + .map(serde_json::Map::len) + .max() + && column_count > MAX_TOOL_COLUMNS + { + return Err(format!( + "query selected {column_count} columns; the tool limit is {MAX_TOOL_COLUMNS}" + )); + } + Ok(value) +} + +fn bounded_select(query: &str) -> Result { + bounded_tool_sql(query) +} + fn visible_datasets(tenant_id: &Option, session_key: &SessionKey) -> Vec { + let internal_datasets = PARSEABLE + .streams + .list_internal_streams(tenant_id) + .into_iter() + .collect::>(); let mut datasets = PARSEABLE .streams .list(tenant_id) .into_iter() + .filter(|dataset| !internal_datasets.contains(dataset)) .filter(|dataset| { Users.authorize(session_key.clone(), Action::ListStream, Some(dataset), None) == Response::Authorized @@ -76,54 +181,107 @@ fn visible_datasets(tenant_id: &Option, session_key: &SessionKey) -> Vec datasets } +async fn visible_dataset_catalog( + tenant_id: &Option, + session_key: &SessionKey, +) -> Vec { + let mut catalog = Vec::new(); + for name in visible_datasets(tenant_id, session_key) { + let telemetry_type = if PARSEABLE.check_or_load_stream(&name, tenant_id).await { + PARSEABLE + .get_stream(&name, tenant_id) + .ok() + .and_then(|stream| { + stream + .metadata + .read() + .ok() + .map(|metadata| metadata.telemetry_type.to_string()) + }) + } else { + None + }; + catalog.push(json!({ + "name": name, + "telemetryType": telemetry_type, + })); + } + catalog +} + fn resolve_time_range(arguments: &Value) -> Result { + resolve_time_range_at(arguments, Utc::now()) +} + +pub fn resolve_time_range_at(arguments: &Value, reference: DateTime) -> Result { let expression = required_string(arguments, "expression")?; - let timezone = arguments + let timezone_name = arguments .get("timezone") .and_then(Value::as_str) .unwrap_or("UTC"); - if !timezone.eq_ignore_ascii_case("utc") { - return Err("OSS time resolution currently requires timezone `UTC`".to_owned()); - } - let now = Utc::now(); + let timezone = timezone_name + .parse::() + .map_err(|_| format!("invalid IANA timezone `{timezone_name}`"))?; let normalized = expression.trim().to_ascii_lowercase(); - let (start, end) = match normalized.as_str() { - "today" => ( - now.date_naive() - .and_hms_opt(0, 0, 0) - .expect("midnight is valid") - .and_utc(), - now, - ), - "yesterday" => { - let today = now - .date_naive() + let range = if normalized == "today" || normalized == "yesterday" { + let date = reference.with_timezone(&timezone).date_naive() + - Duration::days(i64::from(normalized == "yesterday")); + let local_start = date + .and_hms_opt(0, 0, 0) + .ok_or_else(|| "could not construct local day boundary".to_owned())?; + let start = match timezone.from_local_datetime(&local_start) { + LocalResult::Single(start) | LocalResult::Ambiguous(start, _) => { + start.with_timezone(&Utc) + } + LocalResult::None => return Err("local day boundary does not exist".to_owned()), + }; + let end = if normalized == "today" { + reference + } else { + let next_local_start = (date + Duration::days(1)) .and_hms_opt(0, 0, 0) - .expect("midnight is valid") - .and_utc(); - (today - Duration::days(1), today) - } - _ => { - let duration = normalized - .strip_prefix("last ") - .or_else(|| normalized.strip_prefix("past ")) - .unwrap_or(&normalized); - let duration = if duration.chars().any(|character| character.is_ascii_digit()) { - duration.to_owned() - } else { - format!("1 {duration}") - }; - let duration = humantime::parse_duration(&duration) - .map_err(|error| format!("invalid time expression `{expression}`: {error}"))?; - let duration = Duration::from_std(duration).map_err(|error| error.to_string())?; - (now - duration, now) - } + .ok_or_else(|| "could not construct local day boundary".to_owned())?; + match timezone.from_local_datetime(&next_local_start) { + LocalResult::Single(end) | LocalResult::Ambiguous(end, _) => { + end.with_timezone(&Utc) + } + LocalResult::None => return Err("local day boundary does not exist".to_owned()), + } + }; + TimeRange::new(start, end) + } else if let Some(duration) = normalized + .strip_prefix("last ") + .or_else(|| normalized.strip_prefix("past ")) + .or_else(|| normalized.strip_prefix("now-")) + { + let duration = duration.trim(); + let duration_with_quantity = + (!duration.chars().any(|character| character.is_ascii_digit())).then(|| { + let unit = duration + .strip_prefix("an ") + .or_else(|| duration.strip_prefix("a ")) + .unwrap_or(duration); + format!("1 {unit}") + }); + let duration = duration_with_quantity.as_deref().unwrap_or(duration); + let duration = humantime::parse_duration(duration) + .map_err(|error| format!("invalid time duration: {error}"))?; + let duration = Duration::from_std(duration) + .map_err(|error| format!("time duration is too large: {error}"))?; + TimeRange::new(reference - duration, reference) + } else if let Some((start, end)) = expression.split_once(" to ") { + TimeRange::parse_human_time(start.trim(), end.trim()).map_err(|error| error.to_string())? + } else if DateTime::parse_from_rfc3339(expression).is_ok() { + TimeRange::parse_human_time(expression, &reference.to_rfc3339()) + .map_err(|error| error.to_string())? + } else { + TimeRange::parse_human_time(expression, "now").map_err(|error| error.to_string())? }; Ok(json!({ "expression": expression, - "timezone": "UTC", - "startTime": start.to_rfc3339(), - "endTime": end.to_rfc3339(), + "timezone": timezone_name, + "startTime": range.start.to_rfc3339(), + "endTime": range.end.to_rfc3339(), })) } @@ -174,14 +332,177 @@ fn authorize( } } +fn dataset_matches_telemetry_type(dataset: &Value, expected: Option<&str>) -> bool { + expected.is_none_or(|expected| { + dataset + .get("telemetryType") + .and_then(Value::as_str) + .is_some_and(|actual| actual.eq_ignore_ascii_case(expected)) + }) +} + +pub fn resolve_dataset_candidates(payload: Value, arguments: &Value) -> Value { + let query = arguments + .get("query") + .and_then(Value::as_str) + .unwrap_or_default() + .trim() + .to_ascii_lowercase(); + let limit = bounded_integer(arguments, "limit", 5, 1, 20) as usize; + let Some(datasets) = payload.as_array() else { + return payload; + }; + let telemetry_type = arguments + .get("telemetryType") + .and_then(Value::as_str) + .map(str::to_ascii_lowercase); + + let requested_dataset = arguments + .get("selectedDataset") + .or_else(|| arguments.get("contextDataset")) + .and_then(Value::as_str) + .filter(|name| !name.is_empty()); + if let Some(requested_dataset) = requested_dataset { + let selected = datasets.iter().find(|dataset| { + dataset + .get("name") + .and_then(Value::as_str) + .is_some_and(|name| name.eq_ignore_ascii_case(requested_dataset)) + }); + return match selected { + Some(dataset) if dataset_matches_telemetry_type(dataset, telemetry_type.as_deref()) => { + json!({ + "status": "resolved", + "datasets": [dataset], + "selectedDatasets": [dataset.get("name").and_then(Value::as_str).unwrap_or(requested_dataset)], + "confidence": 1.0, + "reason": if arguments.get("selectedDataset").is_some() { + "explicit candidate selection validated against caller-visible datasets" + } else { + "trusted page-context dataset validated against caller-visible datasets" + } + }) + } + Some(_) => json!({ + "status": "not_found", + "datasets": [], + "message": format!( + "Dataset `{requested_dataset}` does not match requested telemetry type `{}`", + telemetry_type.as_deref().unwrap_or_default() + ), + }), + None => json!({ + "status": "error", + "message": format!("Dataset `{requested_dataset}` is not visible to this caller"), + }), + }; + } + + if datasets.len() == 1 { + let dataset = &datasets[0]; + if !dataset_matches_telemetry_type(dataset, telemetry_type.as_deref()) { + return json!({ + "status": "not_found", + "datasets": [], + "message": format!( + "The only caller-visible dataset does not match requested telemetry type `{}`", + telemetry_type.as_deref().unwrap_or_default() + ), + }); + } + return json!({ + "status": "resolved", + "datasets": [dataset], + "selectedDatasets": [dataset.get("name").and_then(Value::as_str).unwrap_or_default()], + "confidence": 1.0, + "reason": "only caller-visible dataset" + }); + } + + let terms = query + .split(|character: char| !character.is_ascii_alphanumeric() && character != '_') + .filter(|term| term.len() > 2) + .collect::>(); + let mut scored = datasets + .iter() + .filter_map(|dataset| { + let original_name = dataset.get("name")?.as_str()?; + if !dataset_matches_telemetry_type(dataset, telemetry_type.as_deref()) { + return None; + } + let name = original_name.to_ascii_lowercase(); + let mut score = 0_u64; + let mut reasons = Vec::new(); + if !query.is_empty() && name == query { + score += 1_000; + reasons.push("exact dataset name"); + } else if !query.is_empty() && query.contains(&name) { + score += 750; + reasons.push("dataset name explicitly mentioned in the question"); + } else if !query.is_empty() && name.contains(&query) { + score += 500; + reasons.push("query matches the dataset name"); + } + let term_matches = terms.iter().filter(|term| name.contains(**term)).count() as u64; + if term_matches > 0 { + score += term_matches * 25; + reasons.push("dataset-name token match"); + } + if telemetry_type.is_some() { + score += 15; + reasons.push("telemetry type match"); + } + Some((score, original_name.to_owned(), reasons, dataset.clone())) + }) + .collect::>(); + scored.sort_by(|left, right| right.0.cmp(&left.0).then_with(|| left.1.cmp(&right.1))); + + if let Some((score, name, _, dataset)) = scored.first() + && *score >= 500 + && scored.get(1).is_none_or(|second| *score > second.0) + { + return json!({ + "status": "resolved", + "datasets": [dataset], + "selectedDatasets": [name], + "confidence": 0.95, + "reason": "unique exact dataset-name match" + }); + } + + let candidates = scored + .into_iter() + .take(limit) + .map(|(score, name, reasons, dataset)| { + json!({ + "name": name, + "score": score, + "reasons": reasons, + "dataset": dataset, + }) + }) + .collect::>(); + json!({ + "status": if candidates.is_empty() { "not_found" } else { "ambiguous" }, + "datasets": [], + "candidates": candidates, + "query": query, + "message": if datasets.is_empty() { + "No caller-visible datasets were found." + } else { + "Dataset resolution is ambiguous. Inspect candidate schemas, then call resolve_dataset again with selectedDataset." + } + }) +} + #[derive(Clone, Copy)] struct OssToolContext<'a> { session_key: &'a SessionKey, tenant_id: &'a Option, - request_headers: &'a HeaderMap, + query_headers: &'a HeaderMap, } -fn execute_dataset_resolution_tool( +async fn execute_dataset_resolution_tool( name: &str, arguments: &Value, context: OssToolContext<'_>, @@ -192,78 +513,169 @@ fn execute_dataset_resolution_tool( .map(|name| json!({ "name": name })) .collect()), "resolve_dataset" => { - let datasets = visible_datasets(context.tenant_id, context.session_key); - let selected = arguments - .get("selectedDataset") - .or_else(|| arguments.get("contextDataset")) - .and_then(Value::as_str) - .or_else(|| arguments.get("query").and_then(Value::as_str)); - if let Some(selected) = selected - && let Some(dataset) = datasets.iter().find(|dataset| dataset.as_str() == selected) - { - return Ok(json!({ - "status": "resolved", - "datasets": [{ "name": dataset }], - "selectedDatasets": [dataset], - "confidence": 1.0, - "reason": "exact caller-visible dataset name", - })); - } - Ok(json!({ - "status": if datasets.is_empty() { "not_found" } else { "ambiguous" }, - "datasets": [], - "candidates": datasets.into_iter().map(|name| json!({ "name": name })).collect::>(), - "message": "Select one exact caller-visible dataset.", - })) + let datasets = visible_dataset_catalog(context.tenant_id, context.session_key).await; + Ok(resolve_dataset_candidates( + Value::Array(datasets), + arguments, + )) } "resolve_time_range" => resolve_time_range(arguments), _ => unreachable!("dataset resolution dispatcher only receives known tools"), } } +enum DatasetActivity { + Skipped, + Count { name: String, events: u64 }, + Error { name: String, message: String }, +} + async fn discover_datasets( arguments: &Value, context: OssToolContext<'_>, ) -> Result { - let telemetry_type = required_string(arguments, "telemetryType")?; - if !matches!(telemetry_type, "logs" | "metrics" | "traces") { + let start_time = required_string(arguments, "startTime")?.to_owned(); + let end_time = required_string(arguments, "endTime")?.to_owned(); + TimeRange::parse_human_time(&start_time, &end_time).map_err(|error| error.to_string())?; + let telemetry_type = required_string(arguments, "telemetryType")?.to_ascii_lowercase(); + if !matches!(telemetry_type.as_str(), "logs" | "metrics" | "traces") { return Err("telemetryType must be logs, metrics, or traces".to_owned()); } - let start = DateTime::parse_from_rfc3339(required_string(arguments, "startTime")?) - .map_err(|error| error.to_string())?; - let end = DateTime::parse_from_rfc3339(required_string(arguments, "endTime")?) - .map_err(|error| error.to_string())?; let limit = bounded_integer(arguments, "limit", 50, 1, 500) as usize; - let mut datasets = Vec::new(); - let mut errors = Vec::new(); - for dataset in visible_datasets(context.tenant_id, context.session_key) { - match get_prism_logstream_info(&dataset, context.tenant_id).await { - Ok(info) => { - let info = serde_json::to_value(info).map_err(|error| error.to_string())?; - let matches_type = nested_string(&info, "telemetryType") - .is_some_and(|value| value.eq_ignore_ascii_case(telemetry_type)); - let active = nested_string(&info, "latestEventAt") - .and_then(|value| DateTime::parse_from_rfc3339(value).ok()) - .is_some_and(|latest| latest >= start && latest <= end); - if matches_type && active { - datasets.push(json!({ "name": dataset, "info": info })); - if datasets.len() == limit { - break; + let datasets = visible_datasets(context.tenant_id, context.session_key); + let visible_datasets = datasets.len(); + + let activity = stream::iter(datasets) + .map(|name| { + let start_time = start_time.clone(); + let end_time = end_time.clone(); + let telemetry_type = telemetry_type.clone(); + async move { + if Users.authorize( + context.session_key.clone(), + Action::Query, + Some(&name), + None, + ) != Response::Authorized + { + return DatasetActivity::Skipped; + } + if !PARSEABLE + .check_or_load_stream(&name, context.tenant_id) + .await + { + return DatasetActivity::Error { + name, + message: "dataset metadata is unavailable".to_owned(), + }; + } + let actual_telemetry_type = match PARSEABLE.get_stream(&name, context.tenant_id) { + Ok(stream) => match stream.metadata.read() { + Ok(metadata) => metadata.telemetry_type.to_string(), + Err(error) => { + return DatasetActivity::Error { + name, + message: error.to_string(), + }; + } + }, + Err(error) => { + return DatasetActivity::Error { + name, + message: error.to_string(), + }; } + }; + if actual_telemetry_type != telemetry_type { + return DatasetActivity::Skipped; + } + + let request = crate::query::CountsRequest { + stream: name.clone(), + start_time, + end_time, + num_bins: Some(1), + conditions: None, + }; + match request.get_bin_density(context.tenant_id).await { + Ok(counts) => DatasetActivity::Count { + name, + events: counts.into_iter().map(|record| record.count).sum(), + }, + Err(error) => DatasetActivity::Error { + name, + message: error.to_string(), + }, } } - Err(error) => errors.push(json!({ - "dataset": dataset, - "message": error.to_string(), - })), + }) + .buffer_unordered(DATASET_DISCOVERY_CONCURRENCY) + .collect::>() + .await; + + let mut active = Vec::new(); + let mut queried_datasets = 0_usize; + let mut errors = Vec::new(); + for result in activity { + match result { + DatasetActivity::Skipped => {} + DatasetActivity::Count { name, events } => { + queried_datasets += 1; + if events > 0 { + active.push((name, events)); + } + } + DatasetActivity::Error { name, message } => { + errors.push(json!({ "name": name, "message": message })); + } } } + active.sort_by(|left, right| right.1.cmp(&left.1).then_with(|| left.0.cmp(&right.0))); + let unavailable_datasets = errors.len(); + + if active.is_empty() && !errors.is_empty() { + tracing::warn!( + telemetry_type, + visible_datasets, + queried_datasets, + unavailable_datasets, + "dataset activity discovery was incomplete" + ); + return Ok(json!({ + "status": "error", + "message": "Dataset activity could not be determined reliably for the requested telemetry type and time range.", + "telemetryType": telemetry_type, + "startTime": start_time, + "endTime": end_time, + "visibleDatasets": visible_datasets, + "queriedDatasets": queried_datasets, + "unavailableDatasets": errors.into_iter().take(20).collect::>(), + })); + } + + let active_count = active.len(); + let selected = active + .into_iter() + .take(limit) + .map(|(name, events)| json!({ "name": name, "eventsInRange": events })) + .collect::>(); + let selected_names = selected + .iter() + .filter_map(|dataset| dataset.get("name").and_then(Value::as_str)) + .collect::>(); Ok(json!({ + "status": if selected.is_empty() { "not_found" } else { "resolved" }, + "selectionSource": "activity_in_resolved_time_range", "telemetryType": telemetry_type, - "startTime": start.to_rfc3339(), - "endTime": end.to_rfc3339(), - "datasets": datasets, - "errors": errors, + "startTime": start_time, + "endTime": end_time, + "queriedDatasets": queried_datasets, + "visibleDatasets": visible_datasets, + "activeDatasetCount": active_count, + "truncated": active_count > selected.len(), + "datasets": selected, + "selectedDatasets": selected_names, + "unavailableDatasets": errors.into_iter().take(20).collect::>(), })) } @@ -318,110 +730,206 @@ async fn execute_dataset_metadata_tool( } } -async fn sample_events(arguments: &Value, context: OssToolContext<'_>) -> Result { +pub fn sample_events_query( + arguments: &Value, + fields: &[String], + reference: DateTime, +) -> Result { let dataset = required_string(arguments, "dataset")?; - authorize(context.session_key, Action::Query, Some(dataset), None)?; let limit = bounded_integer(arguments, "limit", 10, 1, 100); - let minutes = bounded_integer(arguments, "minutes", 60, 1, 1440); - let schema = logstream::get_schema_internal(dataset, context.tenant_id) - .await - .map_err(|error| error.to_string())?; - let fields = schema - .fields() + let (start, end) = match ( + arguments.get("startTime").and_then(Value::as_str), + arguments.get("endTime").and_then(Value::as_str), + ) { + (Some(start), Some(end)) => { + let range = + TimeRange::parse_human_time(start, end).map_err(|error| error.to_string())?; + (range.start, range.end) + } + (None, None) => { + let minutes = bounded_integer(arguments, "minutes", 60, 1, 1440); + (reference - Duration::minutes(minutes as i64), reference) + } + _ => return Err("startTime and endTime must be provided together".to_owned()), + }; + let preferred = [ + "p_timestamp", + "service.name", + "service_name", + "severity_text", + "severity_number", + "body", + "message", + "trace_id", + "span_id", + "duration_nano", + ]; + let mut selected = preferred .iter() - .take(10) - .map(|field| quote_identifier(field.name())) + .filter(|candidate| fields.iter().any(|field| field == **candidate)) + .map(|field| (*field).to_owned()) .collect::>(); - if fields.is_empty() { + for field in fields { + if selected.len() >= MAX_TOOL_COLUMNS { + break; + } + if !selected.contains(field) { + selected.push(field.clone()); + } + } + if selected.is_empty() { return Err("dataset schema has no fields".to_owned()); } - let end = Utc::now(); + let order = if selected.iter().any(|field| field == "p_timestamp") { + " ORDER BY \"p_timestamp\" DESC" + } else { + "" + }; + Ok(query::Query { + query: format!( + "SELECT {} FROM {}{order} LIMIT {limit}", + selected + .iter() + .map(|field| quote_identifier(field)) + .collect::>() + .join(", "), + quote_identifier(dataset), + ), + start_time: start.to_rfc3339(), + end_time: end.to_rfc3339(), + send_null: true, + fields: false, + streaming: false, + filter_tags: None, + }) +} + +pub async fn dataset_field_names( + dataset: &str, + tenant_id: &Option, +) -> Result, String> { + let schema = logstream::get_schema_internal(dataset, tenant_id) + .await + .map_err(|error| error.to_string())?; + Ok(schema + .fields() + .iter() + .map(|field| field.name().to_owned()) + .collect()) +} + +async fn sample_events(arguments: &Value, context: OssToolContext<'_>) -> Result { + let dataset = required_string(arguments, "dataset")?; + authorize(context.session_key, Action::Query, Some(dataset), None)?; + let fields = dataset_field_names(dataset, context.tenant_id).await?; let response = query::query_internal( - query::Query { - query: format!( - "SELECT {} FROM {} LIMIT {limit}", - fields.join(", "), - quote_identifier(dataset) - ), - start_time: (end - Duration::minutes(minutes as i64)).to_rfc3339(), - end_time: end.to_rfc3339(), - send_null: true, - fields: false, - streaming: false, - filter_tags: None, - }, + sample_events_query(arguments, &fields, Utc::now())?, context.session_key, context.tenant_id, ) .await .map_err(|error| error.to_string())?; - response_value(response).await + enforce_tool_tabular_limits(response_value(response).await?) } -fn normalized_log_context_bounds( - arguments: &Value, - anchor: DateTime, -) -> Result<(String, String), String> { - let start_time = arguments - .get("contextStartTime") - .and_then(Value::as_str) - .map(str::to_owned) - .unwrap_or_else(|| (anchor - Duration::minutes(5)).to_rfc3339()); - let end_time = arguments - .get("contextEndTime") - .and_then(Value::as_str) - .map(str::to_owned) - .unwrap_or_else(|| (anchor + Duration::minutes(5)).to_rfc3339()); - let time_range = TimeRange::parse_human_time(&start_time, &end_time) - .map_err(|error| format!("invalid log context time range: {error}"))?; - Ok((time_range.start.to_rfc3339(), time_range.end.to_rfc3339())) +pub fn log_context_request(arguments: &Value) -> Result { + let dataset = required_string(arguments, "dataset")?; + let timestamp = required_string(arguments, "pTimestamp")?; + let mut body = serde_json::Map::from_iter([ + ("dataset".to_owned(), json!(dataset)), + ("pTimestamp".to_owned(), json!(timestamp)), + ]); + let mut has_match_field = false; + for name in ["log", "body", "message"] { + if let Some(value) = arguments.get(name).and_then(Value::as_str) { + body.insert(name.to_owned(), json!(value)); + has_match_field = true; + } + } + if !has_match_field { + return Err( + "get_log_context requires at least one exact `log`, `body`, or `message` value from the anchor row" + .to_owned(), + ); + } + for name in ["contextWindow", "contextStartTime", "contextEndTime"] { + if let Some(value) = arguments.get(name).and_then(Value::as_str) { + body.insert(name.to_owned(), json!(value)); + } + } + if !body.contains_key("contextWindow") + && !body.contains_key("contextStartTime") + && !body.contains_key("contextEndTime") + { + body.insert("contextWindow".to_owned(), json!("5 minutes")); + } + if let Some(conditions) = arguments.get("conditions") { + body.insert("conditions".to_owned(), conditions.clone()); + } + body.insert( + "pageSize".to_owned(), + json!(bounded_integer(arguments, "pageSize", 100, 1, 500)), + ); + Ok(Value::Object(body)) +} + +pub fn project_log_context_result(mut value: Value) -> Value { + const PREFERRED_FIELDS: [&str; 10] = [ + "p_timestamp", + "service.name", + "service_name", + "severity_text", + "severity_description", + "body", + "message", + "log", + "trace_id", + "span_id", + ]; + let rows = match &mut value { + Value::Array(rows) => Some(rows), + Value::Object(object) => object.get_mut("records").and_then(Value::as_array_mut), + _ => None, + }; + if let Some(rows) = rows { + rows.truncate(MAX_TOOL_ROWS as usize); + for row in rows { + let Some(object) = row.as_object_mut() else { + continue; + }; + let mut selected = PREFERRED_FIELDS + .iter() + .filter(|field| object.contains_key(**field)) + .take(MAX_TOOL_COLUMNS) + .map(|field| (*field).to_owned()) + .collect::>(); + for field in object.keys() { + if selected.len() >= MAX_TOOL_COLUMNS { + break; + } + if !selected.contains(field) { + selected.push(field.clone()); + } + } + object.retain(|field, _| selected.contains(field)); + } + } + value } async fn get_log_context(arguments: &Value, context: OssToolContext<'_>) -> Result { + let request = serde_json::from_value::(log_context_request(arguments)?) + .map_err(|error| format!("invalid log context request: {error}"))?; let dataset = required_string(arguments, "dataset")?; authorize(context.session_key, Action::Query, Some(dataset), None)?; - let anchor = DateTime::parse_from_rfc3339(required_string(arguments, "pTimestamp")?) - .map_err(|error| error.to_string())? - .with_timezone(&Utc); - let (start_time, end_time) = normalized_log_context_bounds(arguments, anchor)?; - let limit = bounded_integer(arguments, "pageSize", 100, 1, 500); - let schema = logstream::get_schema_internal(dataset, context.tenant_id) + query::create_streams_for_distributed(vec![dataset.to_owned()], context.tenant_id) .await .map_err(|error| error.to_string())?; - let fields = schema - .fields() - .iter() - .take(10) - .map(|field| quote_identifier(field.name())) - .collect::>(); - if fields.is_empty() { - return Err("dataset schema has no fields".to_owned()); - } - let response = query::query_internal( - query::Query { - query: format!( - "SELECT {} FROM {} WHERE {} >= {} AND {} <= {} ORDER BY {} LIMIT {limit}", - fields.join(", "), - quote_identifier(dataset), - quote_identifier("p_timestamp"), - quote_literal(&start_time), - quote_identifier("p_timestamp"), - quote_literal(&end_time), - quote_identifier("p_timestamp"), - ), - start_time, - end_time, - send_null: true, - fields: false, - streaming: false, - filter_tags: None, - }, - context.session_key, - context.tenant_id, - ) - .await - .map_err(|error| error.to_string())?; - response_value(response).await + let response = query_context::query_context_for_authorized_dataset(request, context.tenant_id) + .await + .map_err(|error| error.to_string())?; + let response = serde_json::to_value(response).map_err(|error| error.to_string())?; + Ok(project_log_context_result(response)) } // Keep the explicit severity policy aligned with frontend/src/lib/errorAnalytics.ts. @@ -799,7 +1307,7 @@ async fn execute_trace_tool( }, context.tenant_id, context.session_key, - Some(context.request_headers.clone()), + Some(context.query_headers.clone()), ) .await } else { @@ -821,7 +1329,7 @@ async fn execute_trace_tool( }, context.tenant_id, context.session_key, - Some(context.request_headers.clone()), + Some(context.query_headers.clone()), ) .await }; @@ -839,6 +1347,189 @@ async fn execute_trace_tool( Ok(json!({ "results": results, "errors": errors })) } +fn time_bucket_interval(start: &str, end: &str, bucket_count: u64) -> String { + let seconds = TimeRange::parse_human_time(start, end) + .map(|range| (range.end - range.start).num_seconds().max(1) as u64) + .unwrap_or(3_600); + format!("{} seconds", (seconds / bucket_count.max(1)).max(1)) +} + +fn trace_activity_query( + dataset: &str, + time_field: &str, + interval: &str, + group_by: &[String], + clauses: &[String], +) -> String { + let bucket = format!( + "CAST(DATE_BIN({}, {}, TIMESTAMP '1970-01-01T00:00:00Z') AS TEXT)", + quote_literal(interval), + quote_identifier(time_field), + ); + let mut select = vec![format!("{bucket} AS \"time\"")]; + select.extend(group_by.iter().map(|field| quote_identifier(field))); + select.push("COUNT(DISTINCT \"span_span_id\") AS \"span_count\"".to_owned()); + select.push("COUNT(DISTINCT \"span_trace_id\") AS \"trace_count\"".to_owned()); + let groups = (1..=group_by.len() + 1) + .map(|position| position.to_string()) + .collect::>() + .join(", "); + let where_clause = clauses + .first() + .map(|_| format!(" WHERE {}", clauses.join(" AND "))) + .unwrap_or_default(); + format!( + "SELECT {} FROM {}{} GROUP BY {} ORDER BY 1 LIMIT 500", + select.join(", "), + quote_identifier(dataset), + where_clause, + groups, + ) +} + +async fn trace_activity_fields( + dataset: &str, + context: OssToolContext<'_>, +) -> Result<(Vec, String), String> { + authorize(context.session_key, Action::Query, Some(dataset), None)?; + if !PARSEABLE + .check_or_load_stream(dataset, context.tenant_id) + .await + { + return Err(format!("dataset `{dataset}` metadata is unavailable")); + } + let stream = PARSEABLE + .get_stream(dataset, context.tenant_id) + .map_err(|error| error.to_string())?; + let telemetry_type = stream + .metadata + .read() + .map_err(|error| error.to_string())? + .telemetry_type + .to_string(); + if telemetry_type != "traces" { + return Err(format!( + "dataset `{dataset}` has telemetry type `{telemetry_type}`; expected `traces`" + )); + } + let schema = logstream::get_schema_internal(dataset, context.tenant_id) + .await + .map_err(|error| error.to_string())?; + let fields = schema + .fields() + .iter() + .map(|field| field.name().to_owned()) + .collect::>(); + for required in ["span_span_id", "span_trace_id"] { + if !fields.iter().any(|field| field == required) { + return Err(format!( + "trace dataset `{dataset}` is missing required field `{required}`" + )); + } + } + let time_field = stream + .get_time_partition() + .unwrap_or_else(|| "p_timestamp".to_owned()); + if !fields.contains(&time_field) { + return Err(format!( + "trace dataset `{dataset}` is missing time field `{time_field}`" + )); + } + Ok((fields, time_field)) +} + +fn trace_activity_group_by(arguments: &Value, fields: &[String]) -> Result, String> { + let group_by = arguments + .get("groupBy") + .and_then(Value::as_array) + .into_iter() + .flatten() + .map(|value| { + value + .as_str() + .map(str::to_owned) + .ok_or_else(|| "groupBy entries must be strings".to_owned()) + }) + .collect::, _>>()?; + if group_by.len() > 5 { + return Err("groupBy may contain at most 5 fields".to_owned()); + } + if let Some(field) = group_by.iter().find(|field| !fields.contains(field)) { + return Err(format!("groupBy field `{field}` does not exist")); + } + Ok(group_by) +} + +fn trace_activity_filter_clauses( + arguments: &Value, + fields: &[String], +) -> Result, String> { + arguments + .get("filters") + .and_then(Value::as_object) + .into_iter() + .flatten() + .map(|(field, value)| { + if !fields.iter().any(|candidate| candidate == field) { + return Err(format!("filter field `{field}` does not exist")); + } + match value { + Value::String(value) => Ok(format!( + "{} = {}", + quote_identifier(field), + quote_literal(value) + )), + Value::Number(value) => Ok(format!("{} = {value}", quote_identifier(field))), + Value::Bool(value) => Ok(format!("{} = {value}", quote_identifier(field))), + Value::Null => Ok(format!("{} IS NULL", quote_identifier(field))), + _ => Err("filter values must be JSON scalars".to_owned()), + } + }) + .collect() +} + +async fn execute_trace_activity_tool( + arguments: &Value, + context: OssToolContext<'_>, +) -> Result { + let dataset = required_string(arguments, "dataset")?; + let start_time = required_string(arguments, "startTime")?; + let end_time = required_string(arguments, "endTime")?; + let (fields, time_field) = trace_activity_fields(dataset, context).await?; + let group_by = trace_activity_group_by(arguments, &fields)?; + let clauses = trace_activity_filter_clauses(arguments, &fields)?; + 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); + let response = query::query_internal( + query::Query { + query: bounded_select(&query)?, + start_time: start_time.to_owned(), + end_time: end_time.to_owned(), + send_null: true, + fields: false, + streaming: false, + filter_tags: None, + }, + context.session_key, + context.tenant_id, + ) + .await + .map_err(|error| error.to_string())?; + let rows = response_value(response).await?; + Ok(json!({ + "dataset": dataset, + "queryType": "sql", + "query": query, + "bucketInterval": interval, + "semantics": { + "span_count": "COUNT(DISTINCT span_span_id)", + "trace_count": "COUNT(DISTINCT span_trace_id)" + }, + "rows": rows, + })) +} + async fn execute_identity_tool( name: &str, arguments: &Value, @@ -1009,12 +1700,34 @@ async fn execute_saved_object_tool( "list_dashboards" => { authorize(context.session_key, Action::ListDashboard, None, None)?; let limit = bounded_integer(arguments, "limit", 100, 1, 500) as usize; - Ok( - dashboards::list_dashboards_internal(limit, None, context.tenant_id) + let query = arguments + .get("query") + .and_then(Value::as_str) + .map(str::to_ascii_lowercase); + let fetch_limit = if query.is_some() { 0 } else { limit }; + Ok(Value::Array( + dashboards::list_dashboards_internal(fetch_limit, None, context.tenant_id) .await .map_err(|e| e.to_string())? - .into(), - ) + .into_iter() + .filter(|dashboard| { + query.as_ref().is_none_or(|query| { + dashboard + .get("title") + .and_then(Value::as_str) + .is_some_and(|title| title.to_ascii_lowercase().contains(query)) + || dashboard.get("tags").and_then(Value::as_array).is_some_and( + |tags| { + tags.iter() + .filter_map(Value::as_str) + .any(|tag| tag.to_ascii_lowercase().contains(query)) + }, + ) + }) + }) + .take(limit) + .collect(), + )) } _ => unreachable!("saved object dispatcher only receives known tools"), } @@ -1025,16 +1738,16 @@ async fn execute_oss_tool_inner( arguments: &Value, session_key: &SessionKey, tenant_id: &Option, - request_headers: &HeaderMap, + query_headers: &HeaderMap, ) -> Result { let context = OssToolContext { session_key, tenant_id, - request_headers, + query_headers, }; match name { "list_datasets" | "resolve_dataset" | "resolve_time_range" => { - execute_dataset_resolution_tool(name, arguments, context) + execute_dataset_resolution_tool(name, arguments, context).await } "discover_datasets" => discover_datasets(arguments, context).await, "get_dataset_schema" | "get_dataset_info" | "get_dataset_stats" => { @@ -1052,6 +1765,7 @@ async fn execute_oss_tool_inner( "get_traces" | "get_trace" | "list_trace_errors" => { execute_trace_tool(name, arguments, context).await } + "get_trace_activity" => execute_trace_activity_tool(arguments, context).await, "ping" | "get_cluster_status" | "get_cluster_metrics" @@ -1070,6 +1784,38 @@ async fn execute_oss_tool_inner( } } +/// Execute a registered OSS tool in-process. Enterprise uses this only after +/// applying its own agent policy and action-confirmation checks, so catalog +/// validation and the underlying platform implementation remain canonical. +pub async fn execute_oss_tool( + name: &str, + arguments: &Value, + session_key: &SessionKey, + tenant_id: &Option, + request_headers: &HeaderMap, +) -> Result { + execute_oss_tool_with_query_headers(name, arguments, session_key, tenant_id, request_headers) + .await +} + +/// Execute a registered OSS tool with separate headers for internal query-node +/// requests. Prism callers use this to preserve their server-authenticated +/// routing headers without changing the external MCP execution path. +pub async fn execute_oss_tool_with_query_headers( + name: &str, + arguments: &Value, + session_key: &SessionKey, + tenant_id: &Option, + query_headers: &HeaderMap, +) -> Result { + let spec = oss_tool_specs() + .into_iter() + .find(|tool| tool.name == name) + .ok_or_else(|| format!("unknown OSS tool `{name}`"))?; + spec.validate_arguments(arguments)?; + execute_oss_tool_inner(name, arguments, session_key, tenant_id, query_headers).await +} + /// Generic execution endpoint used by external MCP adapters. pub async fn call_tool_registry( req: HttpRequest, @@ -1148,7 +1894,7 @@ mod tests { use datafusion::sql::sqlparser::{dialect::GenericDialect, parser::Parser}; for source in [ "SELECT\n1 AS value", - "WITH\nx AS (SELECT 1) SELECT * FROM x", + "WITH\nx AS (SELECT 1 AS value) SELECT value FROM x", "SELECT 1 -- trailing comment", ] { let bounded = super::bounded_select(source).unwrap(); @@ -1169,8 +1915,9 @@ mod tests { use serde_json::{Value, json}; use super::{ - call_tool_registry, list_log_errors_query, log_error_query_request, - normalized_log_context_bounds, + bounded_tool_sql, call_tool_registry, list_log_errors_query, log_context_request, + log_error_query_request, project_log_context_result, resolve_dataset_candidates, + resolve_time_range_at, sample_events_query, trace_activity_query, }; fn log_error_test_range() -> crate::utils::time::TimeRange { @@ -1210,6 +1957,320 @@ mod tests { serde_json::from_slice(&writer.into_inner()).unwrap() } + #[actix_web::test] + async fn shared_sql_limits_reject_wildcards_and_excess_columns() { + assert!(bounded_tool_sql("SELECT a, b FROM logs").is_ok()); + assert!(bounded_tool_sql("SELECT * FROM logs").is_err()); + assert!(bounded_tool_sql("SELECT a,b,c,d,e,f,g,h,i,j,k FROM logs").is_err()); + let bounded = bounded_tool_sql("SELECT a FROM logs -- recent events").unwrap(); + assert!(bounded.contains("-- recent events\n) AS")); + } + + #[actix_web::test] + async fn shared_dataset_resolution_does_not_select_an_arbitrary_fallback() { + let datasets = json!([ + { "name": "application-logs" }, + { "name": "otel-metrics" }, + { "name": "otel-traces" } + ]); + let resolved = + resolve_dataset_candidates(datasets.clone(), &json!({ "query": "show otel-metrics" })); + assert_eq!(resolved["selectedDatasets"], json!(["otel-metrics"])); + + let ambiguous = resolve_dataset_candidates( + datasets, + &json!({ "query": "checkout errors", "limit": 2 }), + ); + assert_eq!(ambiguous["status"], "ambiguous"); + assert_eq!(ambiguous["candidates"].as_array().unwrap().len(), 2); + + let typed = resolve_dataset_candidates( + json!([ + { "name": "otel-data", "telemetryType": "logs" }, + { "name": "system-data", "telemetryType": "metrics" } + ]), + &json!({ "query": "data", "telemetryType": "metrics" }), + ); + assert_eq!(typed["selectedDatasets"], json!(["system-data"])); + + let explicit_mismatch = resolve_dataset_candidates( + json!([{ "name": "otel-data", "telemetryType": "logs" }]), + &json!({ + "query": "otel-data", + "selectedDataset": "otel-data", + "telemetryType": "metrics" + }), + ); + assert_eq!(explicit_mismatch["status"], "not_found"); + + let sole_mismatch = resolve_dataset_candidates( + json!([{ "name": "otel-data", "telemetryType": "logs" }]), + &json!({ "query": "data", "telemetryType": "traces" }), + ); + assert_eq!(sole_mismatch["status"], "not_found"); + } + + #[actix_web::test] + async fn shared_time_resolution_honors_timezone_and_reference() { + let reference = DateTime::parse_from_rfc3339("2026-10-01T12:00:00Z") + .unwrap() + .to_utc(); + let resolved = resolve_time_range_at( + &json!({ + "expression": "last 30 minutes", + "timezone": "Asia/Kolkata" + }), + reference, + ) + .unwrap(); + assert_eq!(resolved["startTime"], "2026-10-01T11:30:00+00:00"); + assert_eq!(resolved["endTime"], "2026-10-01T12:00:00+00:00"); + } + + #[actix_web::test] + async fn log_context_projection_prefers_observability_fields() { + let projected = project_log_context_result(json!([{ + "p_timestamp": "2026-10-22T00:00:00Z", + "service.name": "checkout", + "body": "failed", + "unrelated": "drop" + }])); + assert!(projected[0].get("p_timestamp").is_some()); + assert!(projected[0].get("service.name").is_some()); + assert!(projected[0].get("unrelated").is_some()); + } + + #[actix_web::test] + async fn sample_events_prefers_observability_fields_and_explicit_range() { + let query = sample_events_query( + &json!({ + "dataset": "logs", + "startTime": "2026-10-22T00:00:00Z", + "endTime": "2026-10-22T01:00:00Z", + "limit": 5 + }), + &[ + "unused".to_owned(), + "p_timestamp".to_owned(), + "service.name".to_owned(), + "body".to_owned(), + ], + Utc::now(), + ) + .unwrap(); + + assert_eq!(query.start_time, "2026-10-22T00:00:00+00:00"); + assert!( + query + .query + .contains("\"p_timestamp\", \"service.name\", \"body\"") + ); + assert!(query.query.contains("\"unused\"")); + assert!(query.query.ends_with("LIMIT 5")); + } + + #[actix_web::test] + async fn log_context_requires_an_exact_anchor_value() { + assert!( + log_context_request(&json!({ + "dataset": "logs", + "pTimestamp": "2026-10-22T00:00:00Z" + })) + .is_err() + ); + } + + #[actix_web::test] + async fn trace_activity_uses_distinct_span_and_trace_identifiers() { + let query = trace_activity_query( + "astronomy-shop-traces", + "p_timestamp", + "60 seconds", + &["service.name".to_owned()], + &[], + ); + + assert!(query.contains("COUNT(DISTINCT \"span_span_id\") AS \"span_count\"")); + assert!(query.contains("COUNT(DISTINCT \"span_trace_id\") AS \"trace_count\"")); + assert!(query.contains("\"service.name\"")); + assert!(!query.contains("COUNT(*)")); + } + + async fn run_trace_activity_fixture(group_by: &[String], clauses: &[String]) -> Vec { + let context = datafusion::prelude::SessionContext::new(); + let fixture = context + .sql( + r#"SELECT CAST(at AS TIMESTAMP) AS p_timestamp, trace AS span_trace_id, + span AS span_span_id, service AS "service.name", status AS span_status_code + FROM (VALUES ('2026-10-22T00:00:10Z', 't1', 's1', 'checkout', 0), + ('2026-10-22T00:00:20Z', 't1', 's1', 'checkout', 0), + ('2026-10-22T00:00:30Z', 't1', 's2', 'checkout', 2), + ('2026-10-22T00:00:40Z', 't2', 's3', 'cart', 0), + ('2026-10-22T00:01:10Z', 't3', 's4', 'checkout', 2)) + AS f(at, trace, span, service, status)"#, + ) + .await + .unwrap(); + context + .register_table("traces", fixture.into_view()) + .unwrap(); + 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::>()) + .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 + } + + #[actix_web::test] + async fn trace_activity_query_executes_with_distinct_bucketed_counts() { + let rows = run_trace_activity_fixture(&["service.name".to_owned()], &[]).await; + let summary = rows + .iter() + .map(|row| { + ( + row["time"].as_str().unwrap().to_owned(), + row["service.name"].as_str().unwrap().to_owned(), + row["span_count"].as_u64().unwrap(), + row["trace_count"].as_u64().unwrap(), + ) + }) + .collect::>(); + assert_eq!( + summary, + [ + ("2026-10-22T00:00:00".to_owned(), "cart".to_owned(), 1, 1), + ( + "2026-10-22T00:00:00".to_owned(), + "checkout".to_owned(), + 2, + 1 + ), + ( + "2026-10-22T00:01:00".to_owned(), + "checkout".to_owned(), + 1, + 1 + ), + ] + ); + + let fields = ["span_status_code".to_owned()]; + let clauses = super::trace_activity_filter_clauses( + &json!({ "filters": { "span_status_code": 2 } }), + &fields, + ) + .unwrap(); + let rows = run_trace_activity_fixture(&[], &clauses).await; + assert_eq!(rows.len(), 2); + assert!( + rows.iter() + .all(|row| row["span_count"] == 1 && row["trace_count"] == 1) + ); + } + + #[actix_web::test] + async fn trace_activity_rejects_unknown_fields_and_quotes_filter_values() { + let fields = vec!["service.name".to_owned(), "odd\"field".to_owned()]; + assert!( + super::trace_activity_group_by(&json!({ "groupBy": ["missing"] }), &fields).is_err() + ); + assert!(super::trace_activity_group_by(&json!({ "groupBy": [1] }), &fields).is_err()); + assert_eq!( + super::trace_activity_group_by(&json!({}), &fields).unwrap(), + Vec::::new() + ); + assert!( + super::trace_activity_filter_clauses( + &json!({ "filters": { "missing": "x" } }), + &fields + ) + .is_err() + ); + assert!( + super::trace_activity_filter_clauses( + &json!({ "filters": { "service.name": ["x"] } }), + &fields + ) + .is_err() + ); + let clauses = super::trace_activity_filter_clauses( + &json!({ "filters": { "odd\"field": "it's", "service.name": null } }), + &fields, + ) + .unwrap(); + assert!(clauses.contains(&"\"odd\"\"field\" = 'it''s'".to_owned())); + assert!(clauses.contains(&"\"service.name\" IS NULL".to_owned())); + let query = trace_activity_query("traces", "p_timestamp", "1 seconds", &[], &clauses); + assert!(bounded_tool_sql(&query).is_ok()); + } + + #[actix_web::test] + async fn trace_activity_bucket_interval_is_bounded() { + assert_eq!( + super::time_bucket_interval("2026-10-22T00:00:00Z", "2026-10-22T01:00:00Z", 60), + "60 seconds" + ); + assert_eq!( + super::time_bucket_interval("2026-10-22T00:00:00Z", "2026-10-22T00:00:30Z", 100), + "1 seconds" + ); + } + + #[actix_web::test] + async fn time_resolution_uses_local_day_boundaries_and_rejects_unknown_zones() { + let reference = DateTime::parse_from_rfc3339("2026-10-01T12:00:00Z") + .unwrap() + .to_utc(); + let today = resolve_time_range_at( + &json!({ "expression": "today", "timezone": "Asia/Kolkata" }), + reference, + ) + .unwrap(); + assert_eq!(today["startTime"], "2026-09-30T18:30:00+00:00"); + assert_eq!(today["endTime"], "2026-10-01T12:00:00+00:00"); + assert_eq!(today["timezone"], "Asia/Kolkata"); + let yesterday = resolve_time_range_at( + &json!({ "expression": "yesterday", "timezone": "Asia/Kolkata" }), + reference, + ) + .unwrap(); + assert_eq!(yesterday["startTime"], "2026-09-29T18:30:00+00:00"); + assert_eq!(yesterday["endTime"], "2026-09-30T18:30:00+00:00"); + assert!( + resolve_time_range_at( + &json!({ "expression": "today", "timezone": "Mars/Base" }), + reference + ) + .is_err() + ); + } + + #[actix_web::test] + async fn sample_events_requires_both_explicit_range_bounds() { + let fields = ["p_timestamp".to_owned()]; + for partial in [ + json!({ "dataset": "logs", "startTime": "2026-10-22T00:00:00Z" }), + json!({ "dataset": "logs", "endTime": "2026-10-22T00:00:00Z" }), + ] { + assert!(sample_events_query(&partial, &fields, Utc::now()).is_err()); + } + } + #[actix_web::test] async fn log_error_query_preserves_subminute_ranges() { for time_column in ["p_timestamp", "event.time"] { @@ -1472,22 +2533,6 @@ mod tests { } } - #[actix_web::test] - async fn log_context_bounds_normalize_relative_times() { - let (start, end) = normalized_log_context_bounds( - &json!({ - "contextStartTime": "now-15m", - "contextEndTime": "now", - }), - Utc::now(), - ) - .expect("relative bounds must be accepted"); - let start = start.parse::>().expect("valid start time"); - let end = end.parse::>().expect("valid end time"); - - assert_eq!(end - start, Duration::minutes(15)); - } - #[actix_web::test] async fn invocation_endpoint_requires_authentication() { let app = From 0097de1ece7e2ce2cf3a5d8a42acb2ace8961670 Mon Sep 17 00:00:00 2001 From: VibeSWE <338779468+VibeSWE@users.noreply.github.com> Date: Fri, 23 Oct 2026 22:33:39 -0500 Subject: [PATCH 2/5] ci: migrate all workflows off Blacksmith --- .github/actionlint.yaml | 4 -- .github/workflows/blacksmith-testbox.yml | 30 ------------ .github/workflows/build.yaml | 38 +++++++++++++-- .github/workflows/ci.yaml | 4 +- .github/workflows/coverage.yaml | 36 ++++++++++++-- .github/workflows/frontend.yaml | 4 +- .github/workflows/integration-test.yaml | 27 ++++++++--- .github/workflows/release.yml | 61 +++++++++++++++--------- README.md | 4 +- 9 files changed, 129 insertions(+), 79 deletions(-) delete mode 100644 .github/actionlint.yaml delete mode 100644 .github/workflows/blacksmith-testbox.yml diff --git a/.github/actionlint.yaml b/.github/actionlint.yaml deleted file mode 100644 index 055e04ef6..000000000 --- a/.github/actionlint.yaml +++ /dev/null @@ -1,4 +0,0 @@ -self-hosted-runner: - labels: - - blacksmith-8vcpu-ubuntu-2404 - - blacksmith-8vcpu-ubuntu-2404-arm diff --git a/.github/workflows/blacksmith-testbox.yml b/.github/workflows/blacksmith-testbox.yml deleted file mode 100644 index 2527d614d..000000000 --- a/.github/workflows/blacksmith-testbox.yml +++ /dev/null @@ -1,30 +0,0 @@ -name: Blacksmith Testbox -on: - workflow_dispatch: - inputs: - testbox_id: - type: string - description: Testbox session ID - required: true -permissions: - contents: read -jobs: - frontend: - runs-on: blacksmith-8vcpu-ubuntu-2404 - timeout-minutes: 60 - steps: - - name: Begin Testbox - uses: useblacksmith/begin-testbox@v2 - with: - testbox_id: ${{ inputs.testbox_id }} - - uses: actions/checkout@v4 - - uses: actions/setup-node@v4 - with: - node-version: '24' - cache: npm - cache-dependency-path: frontend/package-lock.json - - run: npm ci && npx playwright install --with-deps chromium - working-directory: frontend - - name: Run Testbox - uses: useblacksmith/run-testbox@v2 - if: always() diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml index 6a38f28cd..8e94a79a6 100644 --- a/.github/workflows/build.yaml +++ b/.github/workflows/build.yaml @@ -13,15 +13,15 @@ jobs: build-default: name: Build Default ${{matrix.target}} runs-on: ${{ matrix.os }} - timeout-minutes: 90 + timeout-minutes: 180 strategy: fail-fast: false matrix: include: # Linux builds - both native, matching the release workflow - - os: blacksmith-8vcpu-ubuntu-2404 + - os: ubuntu-24.04 target: x86_64-unknown-linux-gnu - - os: blacksmith-8vcpu-ubuntu-2404-arm + - os: ubuntu-24.04-arm target: aarch64-unknown-linux-gnu # macOS build - aarch64 only (M1/M2) @@ -34,6 +34,20 @@ jobs: steps: - uses: actions/checkout@v4 + - name: Prepare Linux build resources + if: runner.os == 'Linux' + run: | + sudo rm -rf /usr/local/lib/android /usr/share/dotnet /opt/ghc /opt/hostedtoolcache/CodeQL + echo "CARGO_BUILD_JOBS=2" >> "$GITHUB_ENV" + if sudo fallocate -l 16G /swap.build && sudo chmod 600 /swap.build \ + && sudo mkswap /swap.build && sudo swapon /swap.build; then + echo "Added 16G swap" + else + sudo rm -f /swap.build + echo "::warning title=Swap::Could not add swap; continuing without it" + fi + free -h + df -h / - name: Set up Node.js uses: actions/setup-node@v4 @@ -79,10 +93,24 @@ jobs: # Kafka build for x86_64 Linux only build-kafka: name: Build Kafka x86_64-unknown-linux-gnu - runs-on: blacksmith-8vcpu-ubuntu-2404 - timeout-minutes: 90 + runs-on: ubuntu-24.04 + timeout-minutes: 180 steps: - uses: actions/checkout@v4 + - name: Prepare Linux build resources + if: runner.os == 'Linux' + run: | + sudo rm -rf /usr/local/lib/android /usr/share/dotnet /opt/ghc /opt/hostedtoolcache/CodeQL + echo "CARGO_BUILD_JOBS=2" >> "$GITHUB_ENV" + if sudo fallocate -l 16G /swap.build && sudo chmod 600 /swap.build \ + && sudo mkswap /swap.build && sudo swapon /swap.build; then + echo "Added 16G swap" + else + sudo rm -f /swap.build + echo "::warning title=Swap::Could not add swap; continuing without it" + fi + free -h + df -h / - name: Set up Node.js uses: actions/setup-node@v4 diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 11c9df89b..62bc3729c 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -23,7 +23,7 @@ concurrency: jobs: changes: - runs-on: blacksmith-8vcpu-ubuntu-2404 + runs-on: ubuntu-24.04 timeout-minutes: 5 outputs: frontend: ${{ steps.paths.outputs.frontend }} @@ -94,7 +94,7 @@ jobs: name: CI passed if: always() needs: [changes, frontend, rust, integration] - runs-on: blacksmith-8vcpu-ubuntu-2404 + runs-on: ubuntu-24.04 timeout-minutes: 5 steps: - name: Check required results diff --git a/.github/workflows/coverage.yaml b/.github/workflows/coverage.yaml index 2fd4cd490..100e30c19 100644 --- a/.github/workflows/coverage.yaml +++ b/.github/workflows/coverage.yaml @@ -13,10 +13,24 @@ env: jobs: clippy: - runs-on: blacksmith-8vcpu-ubuntu-2404 - timeout-minutes: 45 + runs-on: ubuntu-24.04 + timeout-minutes: 90 steps: - uses: actions/checkout@v4 + - name: Prepare Linux build resources + if: runner.os == 'Linux' + run: | + sudo rm -rf /usr/local/lib/android /usr/share/dotnet /opt/ghc /opt/hostedtoolcache/CodeQL + echo "CARGO_BUILD_JOBS=2" >> "$GITHUB_ENV" + if sudo fallocate -l 16G /swap.build && sudo chmod 600 /swap.build \ + && sudo mkswap /swap.build && sudo swapon /swap.build; then + echo "Added 16G swap" + else + sudo rm -f /swap.build + echo "::warning title=Swap::Could not add swap; continuing without it" + fi + free -h + df -h / - uses: ./.github/actions/setup-rust-ci with: cache-key: rust-clippy @@ -27,10 +41,24 @@ jobs: run: cargo hack clippy --verbose --each-feature --no-dev-deps -- -D warnings tests: - runs-on: blacksmith-8vcpu-ubuntu-2404 - timeout-minutes: 45 + runs-on: ubuntu-24.04 + timeout-minutes: 90 steps: - uses: actions/checkout@v4 + - name: Prepare Linux build resources + if: runner.os == 'Linux' + run: | + sudo rm -rf /usr/local/lib/android /usr/share/dotnet /opt/ghc /opt/hostedtoolcache/CodeQL + echo "CARGO_BUILD_JOBS=2" >> "$GITHUB_ENV" + if sudo fallocate -l 16G /swap.build && sudo chmod 600 /swap.build \ + && sudo mkswap /swap.build && sudo swapon /swap.build; then + echo "Added 16G swap" + else + sudo rm -f /swap.build + echo "::warning title=Swap::Could not add swap; continuing without it" + fi + free -h + df -h / - uses: ./.github/actions/setup-rust-ci with: cache-key: rust-tests-${{ inputs.coverage && 'coverage' || 'plain' }} diff --git a/.github/workflows/frontend.yaml b/.github/workflows/frontend.yaml index af005f31a..a22314f91 100644 --- a/.github/workflows/frontend.yaml +++ b/.github/workflows/frontend.yaml @@ -10,7 +10,7 @@ defaults: jobs: checks: name: Format, build and unit tests - runs-on: blacksmith-8vcpu-ubuntu-2404 + runs-on: ubuntu-24.04 timeout-minutes: 15 steps: - uses: actions/checkout@v4 @@ -37,7 +37,7 @@ jobs: browser: name: ${{ matrix.suite }} - runs-on: blacksmith-8vcpu-ubuntu-2404 + runs-on: ubuntu-24.04 timeout-minutes: 20 # Keep screenshot fonts, Chromium and OS identical to the reviewed baselines. # Update together with @playwright/test in frontend/package-lock.json. diff --git a/.github/workflows/integration-test.yaml b/.github/workflows/integration-test.yaml index dbb4dfa01..123496873 100644 --- a/.github/workflows/integration-test.yaml +++ b/.github/workflows/integration-test.yaml @@ -11,21 +11,34 @@ permissions: jobs: build-image: - runs-on: blacksmith-8vcpu-ubuntu-2404 - timeout-minutes: 30 + runs-on: ubuntu-24.04 + timeout-minutes: 90 steps: - uses: actions/checkout@v4 + - name: Prepare Linux build resources + if: runner.os == 'Linux' + run: | + sudo rm -rf /usr/local/lib/android /usr/share/dotnet /opt/ghc /opt/hostedtoolcache/CodeQL + if sudo fallocate -l 16G /swap.build && sudo chmod 600 /swap.build \ + && sudo mkswap /swap.build && sudo swapon /swap.build; then + echo "Added 16G swap" + else + sudo rm -f /swap.build + echo "::warning title=Swap::Could not add swap; continuing without it" + fi + free -h + df -h / - name: Set up cached Docker builder - uses: useblacksmith/setup-docker-builder@v2 - with: - cache-key: parseable-integration + uses: docker/setup-buildx-action@v3 - name: Build integration image once - uses: useblacksmith/build-push-action@v2 + uses: docker/build-push-action@v6 with: context: . file: Dockerfile.debug platforms: linux/amd64 tags: parseable-ci:${{ github.sha }} + cache-from: type=gha,scope=parseable-integration + cache-to: type=gha,scope=parseable-integration,mode=max outputs: type=docker,dest=${{ runner.temp }}/parseable-ci.tar - uses: actions/upload-artifact@v4 with: @@ -38,7 +51,7 @@ jobs: quest: name: Quest ${{ matrix.deployment }} (${{ inputs.full && 'full' || 'smoke' }}) needs: build-image - runs-on: blacksmith-8vcpu-ubuntu-2404 + runs-on: ubuntu-24.04 timeout-minutes: 40 strategy: fail-fast: false diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 4e880df45..1b244d586 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -7,7 +7,7 @@ # repository, so that job runs only in parseablehq/parseable. # It also creates a checksum file for the artifacts and publishes it to the release page. # It also builds and pushes a Docker image to Docker Hub with the same version tag. -# Linux jobs run on Blacksmith runners. The build-profile input can select the +# Linux jobs run on GitHub-hosted Ubuntu 24.04 runners. The build-profile input can select the # experimental release-thin-lto profile; such runs upload workflow artifacts # only and never publish to the release page or registry. @@ -42,7 +42,7 @@ env: jobs: build-frontend: name: Package frontend - runs-on: blacksmith-8vcpu-ubuntu-2404 + runs-on: ubuntu-24.04 timeout-minutes: 20 permissions: contents: write @@ -100,17 +100,32 @@ jobs: matrix: include: - target: x86_64-unknown-linux-gnu - os: blacksmith-8vcpu-ubuntu-2404 + os: ubuntu-24.04 # A hosted runner lost contact during the cross-compiled LTO build, # so aarch64 builds natively on an Arm runner. - target: aarch64-unknown-linux-gnu - os: blacksmith-8vcpu-ubuntu-2404-arm - timeout-minutes: 90 + os: ubuntu-24.04-arm + timeout-minutes: 180 steps: - uses: actions/checkout@v4 with: ref: ${{ inputs.build-profile == 'release-thin-lto' && inputs.tag || format('refs/tags/{0}', inputs.tag) }} + - name: Prepare Linux build resources + if: runner.os == 'Linux' + run: | + sudo rm -rf /usr/local/lib/android /usr/share/dotnet /opt/ghc /opt/hostedtoolcache/CodeQL + echo "CARGO_BUILD_JOBS=2" >> "$GITHUB_ENV" + if sudo fallocate -l 16G /swap.build && sudo chmod 600 /swap.build \ + && sudo mkswap /swap.build && sudo swapon /swap.build; then + echo "Added 16G swap" + else + sudo rm -f /swap.build + echo "::warning title=Swap::Could not add swap; continuing without it" + fi + free -h + df -h / + - name: Set up Node.js uses: actions/setup-node@v4 with: @@ -140,21 +155,6 @@ jobs: with: key: ${{ matrix.target }}-default-${{ env.RELEASE_PROFILE }} - # The aarch64 runner lost contact during the fat-LTO build twice, which - # GitHub attributes to CPU or memory starvation. Free disk, add swap and - # record the resources as annotations, which survive a lost runner. - # Swap is best effort because runner disk layouts differ. - - name: Free disk and add swap - run: | - sudo rm -rf /usr/local/lib/android /usr/share/dotnet /opt/ghc /opt/hostedtoolcache/CodeQL - if sudo fallocate -l 16G /swap.release && sudo chmod 600 /swap.release \ - && sudo mkswap /swap.release && sudo swapon /swap.release; then - echo "Added 16G swap" - else - echo "::warning title=Swap::Could not add swap; continuing without it" - fi - echo "::notice title=Runner resources::$(nproc) CPUs; $(free -h | awk '/Mem/{print $2" RAM"} /Swap/{print $2" swap"}' | paste -sd ' '); $(df -h --output=avail / | tail -1 | tr -d ' ') free on /" - - name: Build native run: cargo build --profile "$RELEASE_PROFILE" --target ${{ matrix.target }} @@ -332,8 +332,8 @@ jobs: build-kafka-linux: name: Build Kafka for x86_64-unknown-linux-gnu - runs-on: blacksmith-8vcpu-ubuntu-2404 - timeout-minutes: 90 + runs-on: ubuntu-24.04 + timeout-minutes: 180 permissions: id-token: write contents: write @@ -343,6 +343,21 @@ jobs: with: ref: ${{ inputs.build-profile == 'release-thin-lto' && inputs.tag || format('refs/tags/{0}', inputs.tag) }} + - name: Prepare Linux build resources + if: runner.os == 'Linux' + run: | + sudo rm -rf /usr/local/lib/android /usr/share/dotnet /opt/ghc /opt/hostedtoolcache/CodeQL + echo "CARGO_BUILD_JOBS=2" >> "$GITHUB_ENV" + if sudo fallocate -l 16G /swap.build && sudo chmod 600 /swap.build \ + && sudo mkswap /swap.build && sudo swapon /swap.build; then + echo "Added 16G swap" + else + sudo rm -f /swap.build + echo "::warning title=Swap::Could not add swap; continuing without it" + fi + free -h + df -h / + - name: Set up Node.js uses: actions/setup-node@v4 with: @@ -538,7 +553,7 @@ jobs: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} create-checksum: - runs-on: blacksmith-8vcpu-ubuntu-2404 + runs-on: ubuntu-24.04 timeout-minutes: 10 needs: [build-linux, build-windows, build-mac, build-kafka-linux, build-kafka-mac] permissions: diff --git a/README.md b/README.md index 5d0dce79c..b8f015ad2 100644 --- a/README.md +++ b/README.md @@ -80,8 +80,8 @@ For production deployments, refer to the [installation guide ↗︎](https://www ## Contributing -Pull requests use the [CI workflow](.github/workflows/ci.yaml) on Blacksmith -8-vCPU Ubuntu 24.04 runners. It selects frontend, Rust, and integration jobs from +Pull requests use the [CI workflow](.github/workflows/ci.yaml) on GitHub-hosted +Ubuntu 24.04 runners. It selects frontend, Rust, and integration jobs from the changed paths and reports one aggregate `CI passed` check. Frontend changes still exercise the embedded UI through the integration image; Rust-only changes skip the browser suite. A newer PR commit cancels its superseded CI run. From cc46ee5fa2127beabc71919b9152a6db3c9e5d62 Mon Sep 17 00:00:00 2001 From: VibeSWE <338779468+VibeSWE@users.noreply.github.com> Date: Sat, 24 Oct 2026 11:45:37 -0500 Subject: [PATCH 3/5] Enforce JSON Schema type unions in tool argument validation `validate_type` ignored array-valued `type` (for example `["string", "null"]`), so the dashboard operation schemas' nullable and union fields were never type-checked. Accept a value when it matches any listed type, and name every allowed type in the error. Verified the dashboard agent E2E script's create, update, rename, placement, invalid-layout and protected-field payloads still pass schema validation, so service-side 422/400 responses are unchanged. Co-Authored-By: Claude Opus 5.5 --- docs/tool-api.md | 2 +- src/tool_catalog.rs | 31 ++++++++++++++++++++++++++----- 2 files changed, 27 insertions(+), 6 deletions(-) diff --git a/docs/tool-api.md b/docs/tool-api.md index b57ed774d..fc9cca74f 100644 --- a/docs/tool-api.md +++ b/docs/tool-api.md @@ -217,7 +217,7 @@ rerun that SQL unchanged, for example in a dashboard tile. ## Argument schema validation -`inputSchema` is enforced server-side for `type`, `required`, +`inputSchema` is enforced server-side for `type` (a single type or a list of types), `required`, `additionalProperties`, `enum`, `const`, `minLength`/`maxLength`, `minimum`/`maximum`, `minItems`/`maxItems`, `items`, `anyOf` and `oneOf`. `oneOf` passes only when exactly one variant matches. Combinators are checked diff --git a/src/tool_catalog.rs b/src/tool_catalog.rs index 8edc2f833..ac6c955c1 100644 --- a/src/tool_catalog.rs +++ b/src/tool_catalog.rs @@ -194,10 +194,14 @@ fn validate_enum(schema: &Value, value: &Value, path: &str) -> Result<(), String } fn validate_type(schema: &Value, value: &Value, path: &str) -> Result<(), String> { - let Some(expected_type) = schema.get("type").and_then(Value::as_str) else { - return Ok(()); + let expected_types = match schema.get("type") { + Some(Value::String(expected_type)) => vec![expected_type.as_str()], + Some(Value::Array(expected_types)) => { + expected_types.iter().filter_map(Value::as_str).collect() + } + _ => return Ok(()), }; - let type_matches = match expected_type { + let type_matches = |expected_type: &str| match expected_type { "object" => value.is_object(), "array" => value.is_array(), "string" => value.is_string(), @@ -207,10 +211,10 @@ fn validate_type(schema: &Value, value: &Value, path: &str) -> Result<(), String "null" => value.is_null(), _ => true, }; - if type_matches { + if expected_types.is_empty() || expected_types.iter().any(|kind| type_matches(kind)) { Ok(()) } else { - Err(format!("{path} must be a {expected_type}")) + Err(format!("{path} must be a {}", expected_types.join(" or "))) } } @@ -1266,6 +1270,23 @@ mod tests { ); } + #[actix_web::test] + async fn schema_type_unions_accept_any_listed_type() { + let schema = serde_json::json!({ "type": ["string", "null"] }); + for value in [serde_json::json!("x"), serde_json::json!(null)] { + assert!(validate_tool_arguments(&schema, &value).is_ok(), "{value}"); + } + let error = validate_tool_arguments(&schema, &serde_json::json!(1)).unwrap_err(); + assert_eq!(error, "arguments must be a string or null"); + assert!( + validate_tool_arguments( + &serde_json::json!({ "type": "integer" }), + &serde_json::json!(1.5) + ) + .is_err() + ); + } + #[actix_web::test] async fn tool_results_use_mcp_call_tool_result_shape() { let result = serde_json::to_value(ToolCallResult::success(serde_json::json!({ From 621853828b9e4ee72c9b148301d76fbbdeed53fb Mon Sep 17 00:00:00 2001 From: VibeSWE <338779468+VibeSWE@users.noreply.github.com> Date: Sat, 24 Oct 2026 14:25:02 -0500 Subject: [PATCH 4/5] Address PR #62 review: listed tool schemas, log context anchors, overflow, tests - list_tool_registry omits root-level oneOf/anyOf/allOf from each listed inputSchema, since LLM tool APIs reject them; the catalog schema keeps them and call_tool_registry still validates against it. - log_context_request requires exactly one of log, body or message. - resolve_time_range and the log-context window use checked chrono arithmetic, so huge durations return errors instead of panicking. - Extract trace_activity_time_field and summarize_dataset_activity so the trace preconditions and discover_datasets ranking can be unit tested. - Tests: no listed root combinators; exactly-one anchor; oneOf error text; log-context projection with more than 10 fields; log_context_request defaults; sample_events 10-column cap; trace activity ordering with millisecond timestamps and Float64 status; SQL validator query forms; resolve_dataset ties, case-insensitive and invisible selections; duration overflow. Co-Authored-By: Claude Opus 5.5 --- src/handlers/http/query_context.rs | 31 +- src/tool_catalog.rs | 116 +++++++- src/tool_executor.rs | 435 +++++++++++++++++++++++++---- 3 files changed, 525 insertions(+), 57 deletions(-) diff --git a/src/handlers/http/query_context.rs b/src/handlers/http/query_context.rs index 64df0039b..23e768ddd 100644 --- a/src/handlers/http/query_context.rs +++ b/src/handlers/http/query_context.rs @@ -376,11 +376,23 @@ fn log_context_window_bounds( let duration = chrono::Duration::from_std(duration) .map_err(|err| QueryError::CustomError(format!("Invalid contextWindow: {err}")))?; - let start = truncate_to_minute(anchor_timestamp - duration); - let mut end = truncate_to_minute(anchor_timestamp + duration); + let out_of_range = + || QueryError::CustomError("contextWindow is outside the supported range".to_string()); + let start = truncate_to_minute( + anchor_timestamp + .checked_sub_signed(duration) + .ok_or_else(out_of_range)?, + ); + let mut end = truncate_to_minute( + anchor_timestamp + .checked_add_signed(duration) + .ok_or_else(out_of_range)?, + ); if start >= end { - end = start + chrono::Duration::minutes(1); + end = start + .checked_add_signed(chrono::Duration::minutes(1)) + .ok_or_else(out_of_range)?; } Ok((start, end)) @@ -1202,6 +1214,19 @@ mod tests { assert_eq!(format_log_context_api_time(end), "2026-06-17T10:16:00Z"); } + #[test] + fn log_context_window_bounds_reject_windows_beyond_the_supported_range() { + for window in ["300000y", "1000000 years"] { + let error = log_context_window_bounds(anchor_timestamp(), window) + .unwrap_err() + .to_string(); + assert!( + error.contains("contextWindow is outside the supported range"), + "{window}: {error}" + ); + } + } + #[test] fn log_context_bounds_resolver_accepts_one_mode_only() { let anchor = anchor_timestamp(); diff --git a/src/tool_catalog.rs b/src/tool_catalog.rs index ac6c955c1..cb7ecc710 100644 --- a/src/tool_catalog.rs +++ b/src/tool_catalog.rs @@ -853,6 +853,25 @@ pub fn external_tool_specs(writes_enabled: bool) -> impl Iterator Value { + let mut schema = schema.clone(); + if let Some(object) = schema.as_object_mut() { + for combinator in ROOT_SCHEMA_COMBINATORS { + object.remove(combinator); + } + } + schema +} + /// `GET /api/prism/v1/llm/tools/list` pub async fn list_tool_registry(request: HttpRequest) -> HttpResponse { let tools = external_tool_specs(crate::users::dashboard_tools::writes_enabled(&request)) @@ -861,7 +880,7 @@ pub async fn list_tool_registry(request: HttpRequest) -> HttpResponse { "name": tool.name, "title": tool.title, "description": tool.description, - "inputSchema": tool.input_schema, + "inputSchema": listed_input_schema(&tool.input_schema), "annotations": { "readOnlyHint": matches!(tool.effect(), ToolEffect::ReadOnly | ToolEffect::ExpensiveRead), "destructiveHint": tool.name == "delete_dashboard", "idempotentHint": crate::users::dashboard_tools::is_write(tool.name), "openWorldHint": false }, }); if let Some(schema) = crate::users::dashboard_tools::output_schema(tool.name) { value["outputSchema"] = schema; } @@ -881,8 +900,9 @@ mod tests { use serde_json::Value; use super::{ - ToolCallRequest, ToolCallResult, ToolEffect, is_oss_read_tool, list_tool_registry, - oss_action_tool_specs, oss_read_tool_specs, oss_tool_specs, + ROOT_SCHEMA_COMBINATORS, ToolCallRequest, ToolCallResult, ToolEffect, external_tool_specs, + is_oss_read_tool, list_tool_registry, listed_input_schema, oss_action_tool_specs, + oss_read_tool_specs, oss_tool_specs, }; #[actix_web::test] @@ -1074,6 +1094,62 @@ mod tests { assert_eq!(tool["annotations"]["destructiveHint"], false); } + #[actix_web::test] + async fn listed_tools_have_no_root_schema_combinators() { + let app = + test::init_service(App::new().route("/tools/list", web::get().to(list_tool_registry))) + .await; + let response: Value = test::call_and_read_body_json( + &app, + test::TestRequest::get().uri("/tools/list").to_request(), + ) + .await; + let listed = response["tools"] + .as_array() + .unwrap() + .iter() + .map(|tool| { + ( + tool["name"].as_str().unwrap().to_owned(), + tool["inputSchema"].clone(), + ) + }) + .chain(external_tool_specs(true).map(|tool| { + ( + tool.name.to_owned(), + listed_input_schema(&tool.input_schema), + ) + })) + .collect::>(); + assert!(listed.iter().any(|(name, _)| name == "get_log_context")); + for (name, schema) in &listed { + assert_eq!(schema["type"], "object", "{name}"); + for combinator in ROOT_SCHEMA_COMBINATORS { + assert!( + schema.get(combinator).is_none(), + "{name} lists root {combinator}" + ); + } + } + + // The catalog keeps the constraint for server-side validation, and the + // listed property descriptions state it. + let context = external_tool_specs(true) + .find(|tool| tool.name == "get_log_context") + .unwrap(); + assert!(context.input_schema.get("oneOf").is_some()); + let listed_context = listed_input_schema(&context.input_schema); + for anchor in ["log", "body", "message"] { + assert!( + listed_context["properties"][anchor]["description"] + .as_str() + .unwrap() + .contains("exactly one of log, body, or message"), + "{anchor}" + ); + } + } + #[actix_web::test] async fn changed_read_tool_schemas_reject_out_of_contract_arguments() { let find = |name: &str| { @@ -1115,6 +1191,21 @@ mod tests { })) .is_err() ); + for invalid in [ + serde_json::json!({ "dataset": "logs", "pTimestamp": "2026-10-22T00:00:00Z" }), + serde_json::json!({ + "dataset": "logs", + "pTimestamp": "2026-10-22T00:00:00Z", + "log": "a", + "body": "b" + }), + ] { + assert_eq!( + context.validate_arguments(&invalid).unwrap_err(), + "arguments must match exactly one allowed schema", + "{invalid}" + ); + } let sample = find("sample_events"); assert!( @@ -1257,10 +1348,21 @@ mod tests { ] { assert!(validate_tool_arguments(&schema, &value).is_err()); } - assert!( - validate_tool_arguments(&serde_json::json!({"oneOf":[{},{}]}), &serde_json::json!(1)) - .is_err() - ); + for (schema, value) in [ + // Zero matching variants. + ( + serde_json::json!({"oneOf":[{"type":"string"},{"type":"boolean"}]}), + serde_json::json!(1), + ), + // Two matching variants. + (serde_json::json!({"oneOf":[{},{}]}), serde_json::json!(1)), + ] { + assert_eq!( + validate_tool_arguments(&schema, &value).unwrap_err(), + "arguments must match exactly one allowed schema", + "{schema}" + ); + } assert!( validate_tool_arguments( &serde_json::json!({"anyOf":[{"type":"string"}],"maxLength":2}), diff --git a/src/tool_executor.rs b/src/tool_executor.rs index f2c477041..360632335 100644 --- a/src/tool_executor.rs +++ b/src/tool_executor.rs @@ -268,7 +268,10 @@ pub fn resolve_time_range_at(arguments: &Value, reference: DateTime) -> Res .map_err(|error| format!("invalid time duration: {error}"))?; let duration = Duration::from_std(duration) .map_err(|error| format!("time duration is too large: {error}"))?; - TimeRange::new(reference - duration, reference) + let start = reference + .checked_sub_signed(duration) + .ok_or_else(|| "time duration is too large".to_owned())?; + TimeRange::new(start, reference) } else if let Some((start, end)) = expression.split_once(" to ") { TimeRange::parse_human_time(start.trim(), end.trim()).map_err(|error| error.to_string())? } else if DateTime::parse_from_rfc3339(expression).is_ok() { @@ -613,6 +616,25 @@ async fn discover_datasets( .collect::>() .await; + Ok(summarize_dataset_activity( + activity, + limit, + visible_datasets, + &telemetry_type, + &start_time, + &end_time, + )) +} + +/// Builds the `discover_datasets` result from per-dataset activity. +fn summarize_dataset_activity( + activity: Vec, + limit: usize, + visible_datasets: usize, + telemetry_type: &str, + start_time: &str, + end_time: &str, +) -> Value { let mut active = Vec::new(); let mut queried_datasets = 0_usize; let mut errors = Vec::new(); @@ -641,7 +663,7 @@ async fn discover_datasets( unavailable_datasets, "dataset activity discovery was incomplete" ); - return Ok(json!({ + return json!({ "status": "error", "message": "Dataset activity could not be determined reliably for the requested telemetry type and time range.", "telemetryType": telemetry_type, @@ -650,7 +672,7 @@ async fn discover_datasets( "visibleDatasets": visible_datasets, "queriedDatasets": queried_datasets, "unavailableDatasets": errors.into_iter().take(20).collect::>(), - })); + }); } let active_count = active.len(); @@ -663,7 +685,7 @@ async fn discover_datasets( .iter() .filter_map(|dataset| dataset.get("name").and_then(Value::as_str)) .collect::>(); - Ok(json!({ + json!({ "status": if selected.is_empty() { "not_found" } else { "resolved" }, "selectionSource": "activity_in_resolved_time_range", "telemetryType": telemetry_type, @@ -676,7 +698,7 @@ async fn discover_datasets( "datasets": selected, "selectedDatasets": selected_names, "unavailableDatasets": errors.into_iter().take(20).collect::>(), - })) + }) } fn schema_tool_result(schema: &arrow_schema::Schema, arguments: &Value) -> Result { @@ -839,16 +861,16 @@ pub fn log_context_request(arguments: &Value) -> Result { ("dataset".to_owned(), json!(dataset)), ("pTimestamp".to_owned(), json!(timestamp)), ]); - let mut has_match_field = false; + let mut match_fields = 0; for name in ["log", "body", "message"] { if let Some(value) = arguments.get(name).and_then(Value::as_str) { body.insert(name.to_owned(), json!(value)); - has_match_field = true; + match_fields += 1; } } - if !has_match_field { + if match_fields != 1 { return Err( - "get_log_context requires at least one exact `log`, `body`, or `message` value from the anchor row" + "get_log_context requires exactly one exact `log`, `body`, or `message` value from the anchor row" .to_owned(), ); } @@ -1420,6 +1442,17 @@ async fn trace_activity_fields( .iter() .map(|field| field.name().to_owned()) .collect::>(); + let time_field = trace_activity_time_field(dataset, &fields, stream.get_time_partition())?; + Ok((fields, time_field)) +} + +/// Checks that a traces dataset's fields support activity counts and returns +/// the time column to bucket on. +fn trace_activity_time_field( + dataset: &str, + fields: &[String], + time_partition: Option, +) -> Result { for required in ["span_span_id", "span_trace_id"] { if !fields.iter().any(|field| field == required) { return Err(format!( @@ -1427,15 +1460,13 @@ async fn trace_activity_fields( )); } } - let time_field = stream - .get_time_partition() - .unwrap_or_else(|| "p_timestamp".to_owned()); + let time_field = time_partition.unwrap_or_else(|| "p_timestamp".to_owned()); if !fields.contains(&time_field) { return Err(format!( "trace dataset `{dataset}` is missing time field `{time_field}`" )); } - Ok((fields, time_field)) + Ok(time_field) } fn trace_activity_group_by(arguments: &Value, fields: &[String]) -> Result, String> { @@ -1914,6 +1945,8 @@ mod tests { use chrono::{DateTime, Duration, Utc}; use serde_json::{Value, json}; + use crate::handlers::http::query_context::LogContextRequest; + use super::{ bounded_tool_sql, call_tool_registry, list_log_errors_query, log_context_request, log_error_query_request, project_log_context_result, resolve_dataset_candidates, @@ -1966,6 +1999,82 @@ mod tests { assert!(bounded.contains("-- recent events\n) AS")); } + #[actix_web::test] + async fn sql_validator_accepts_and_rejects_query_forms() { + for accepted in [ + "SELECT COUNT(*) FROM t", + "SELECT a FROM t UNION SELECT b FROM u", + "SELECT CAST(a AS BIGINT), TRY_CAST(b AS DOUBLE) FROM t", + ] { + assert!(super::validate_tool_sql(accepted).is_ok(), "{accepted}"); + } + for (rejected, message) in [ + ( + "SELECT a FROM t UNION SELECT * FROM u", + "wildcard selection is not allowed", + ), + ("(SELECT * FROM t)", "wildcard selection is not allowed"), + ("SELECT t.* FROM t", "wildcard selection is not allowed"), + ( + "SELECT a::INT64 FROM t", + "unsupported DataFusion SQL cast type `INT64`; use `BIGINT` instead", + ), + ( + "SELECT TRY_CAST(a AS FLOAT64) FROM t", + "unsupported DataFusion SQL cast type `FLOAT64`; use `DOUBLE` instead", + ), + ( + "SELECT CAST(a AS UINT64) FROM t", + "unsupported DataFusion SQL cast type `UInt64`; use `BIGINT` instead", + ), + ("VALUES (1)", "unsupported SQL query form"), + ( + "EXPLAIN SELECT a FROM t", + "only SQL SELECT queries are allowed", + ), + ("DELETE FROM t", "only SQL SELECT queries are allowed"), + ( + "SELECT a FROM t; SELECT b FROM t", + "exactly one SQL statement is required", + ), + ] { + let error = super::validate_tool_sql(rejected).unwrap_err(); + assert!(error.contains(message), "{rejected}: {error}"); + } + } + + #[actix_web::test] + async fn dataset_resolution_reports_ties_case_insensitive_and_invisible_selections() { + let datasets = json!([ + { "name": "logs" }, + { "name": "app-logs" }, + { "name": "logs-archive" } + ]); + let tied = resolve_dataset_candidates( + datasets.clone(), + &json!({ "query": "app-logs and logs-archive" }), + ); + assert_eq!(tied["status"], "ambiguous"); + assert_eq!(tied["datasets"], json!([])); + assert_eq!( + tied["candidates"][0]["score"], + tied["candidates"][1]["score"] + ); + + let selected = + resolve_dataset_candidates(datasets.clone(), &json!({ "selectedDataset": "APP-LOGS" })); + assert_eq!(selected["status"], "resolved"); + assert_eq!(selected["selectedDatasets"], json!(["app-logs"])); + + let invisible = + resolve_dataset_candidates(datasets, &json!({ "selectedDataset": "secret-logs" })); + assert_eq!(invisible["status"], "error"); + assert_eq!( + invisible["message"], + "Dataset `secret-logs` is not visible to this caller" + ); + } + #[actix_web::test] async fn shared_dataset_resolution_does_not_select_an_arbitrary_fallback() { let datasets = json!([ @@ -2029,15 +2138,68 @@ mod tests { #[actix_web::test] async fn log_context_projection_prefers_observability_fields() { - let projected = project_log_context_result(json!([{ - "p_timestamp": "2026-10-22T00:00:00Z", - "service.name": "checkout", - "body": "failed", - "unrelated": "drop" - }])); - assert!(projected[0].get("p_timestamp").is_some()); - assert!(projected[0].get("service.name").is_some()); - assert!(projected[0].get("unrelated").is_some()); + // Unrelated fields sort before, and are inserted before, the preferred + // ones, so the kept set does not depend on map ordering. + let mut row = serde_json::Map::new(); + for index in 0..12 { + row.insert(format!("a{index:02}"), json!(index)); + } + row.insert("body".to_owned(), json!("failed")); + row.insert("p_timestamp".to_owned(), json!("2026-10-22T00:00:00Z")); + row.insert("service.name".to_owned(), json!("checkout")); + let row = Value::Object(row); + + for (shape, projected) in [ + ("array", project_log_context_result(json!([row.clone()]))), + ( + "records", + project_log_context_result(json!({ "anchorIndex": 0, "records": [row] })) + ["records"] + .clone(), + ), + ] { + let projected = projected[0].as_object().unwrap(); + assert_eq!(projected.len(), 10, "{shape}"); + for kept in ["p_timestamp", "service.name", "body", "a00", "a06"] { + assert!(projected.contains_key(kept), "{shape}: {kept}"); + } + for dropped in ["a07", "a11"] { + assert!(!projected.contains_key(dropped), "{shape}: {dropped}"); + } + } + } + + #[actix_web::test] + async fn log_context_request_defaults_and_bounds() { + let anchor = json!({ + "dataset": "logs", + "pTimestamp": "2026-10-22T00:00:00Z", + "body": "failed" + }); + let request = log_context_request(&anchor).unwrap(); + assert_eq!(request["contextWindow"], "5 minutes"); + assert_eq!(request["pageSize"], 100); + let parsed = serde_json::from_value::(request).unwrap(); + assert_eq!(parsed.context_window.as_deref(), Some("5 minutes")); + assert_eq!(parsed.body.as_deref(), Some("failed")); + + let mut explicit = anchor.clone(); + explicit["contextStartTime"] = json!("2026-10-21T23:50:00Z"); + explicit["contextEndTime"] = json!("2026-10-22T00:10:00Z"); + explicit["pageSize"] = json!(9999); + let request = log_context_request(&explicit).unwrap(); + assert!(request.get("contextWindow").is_none()); + assert_eq!(request["pageSize"], 500); + let parsed = serde_json::from_value::(request).unwrap(); + assert_eq!( + parsed.context_start_time.as_deref(), + Some("2026-10-21T23:50:00Z") + ); + assert_eq!( + parsed.context_end_time.as_deref(), + Some("2026-10-22T00:10:00Z") + ); + assert_eq!(parsed.page_size, Some(500)); } #[actix_web::test] @@ -2070,14 +2232,32 @@ mod tests { } #[actix_web::test] - async fn log_context_requires_an_exact_anchor_value() { - assert!( - log_context_request(&json!({ + async fn log_context_requires_exactly_one_exact_anchor_value() { + for anchors in [ + json!({}), + json!({ "log": "a", "body": "b" }), + json!({ "log": "a", "body": "b", "message": "c" }), + ] { + let mut arguments = json!({ "dataset": "logs", "pTimestamp": "2026-10-22T00:00:00Z" + }); + arguments + .as_object_mut() + .unwrap() + .extend(anchors.as_object().unwrap().clone()); + let error = log_context_request(&arguments).unwrap_err(); + assert!(error.contains("exactly one"), "{anchors}: {error}"); + } + for anchor in ["log", "body", "message"] { + let request = log_context_request(&json!({ + "dataset": "logs", + "pTimestamp": "2026-10-22T00:00:00Z", + anchor: "failed" })) - .is_err() - ); + .unwrap(); + assert_eq!(request[anchor], "failed"); + } } #[actix_web::test] @@ -2096,19 +2276,27 @@ mod tests { assert!(!query.contains("COUNT(*)")); } - async fn run_trace_activity_fixture(group_by: &[String], clauses: &[String]) -> Vec { + /// Runs the trace-activity SQL over a fixture with Parseable's millisecond + /// timestamps, and checks that rows come back ordered by bucket. + async fn run_trace_activity_fixture( + group_by: &[String], + clauses: &[String], + status_type: &str, + ) -> Vec { let context = datafusion::prelude::SessionContext::new(); + // Later buckets come first so that ordering depends on the query. let fixture = context - .sql( - r#"SELECT CAST(at AS TIMESTAMP) AS p_timestamp, trace AS span_trace_id, - span AS span_span_id, service AS "service.name", status AS span_status_code - FROM (VALUES ('2026-10-22T00:00:10Z', 't1', 's1', 'checkout', 0), - ('2026-10-22T00:00:20Z', 't1', 's1', 'checkout', 0), - ('2026-10-22T00:00:30Z', 't1', 's2', 'checkout', 2), + .sql(&format!( + r#"SELECT arrow_cast(at, 'Timestamp(Millisecond, None)') AS p_timestamp, + trace AS span_trace_id, span AS span_span_id, service AS "service.name", + CAST(status AS {status_type}) AS span_status_code + FROM (VALUES ('2026-10-22T00:01:10Z', 't3', 's4', 'checkout', 2), ('2026-10-22T00:00:40Z', 't2', 's3', 'cart', 0), - ('2026-10-22T00:01:10Z', 't3', 's4', 'checkout', 2)) - AS f(at, trace, span, service, status)"#, - ) + ('2026-10-22T00:00:30Z', 't1', 's2', 'checkout', 2), + ('2026-10-22T00:00:20Z', 't1', 's1', 'checkout', 0), + ('2026-10-22T00:00:10Z', 't1', 's1', 'checkout', 0)) + AS f(at, trace, span, service, status)"# + )) .await .unwrap(); context @@ -2131,15 +2319,22 @@ mod tests { .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()); + let rows: Vec = serde_json::from_slice(&writer.into_inner()).unwrap(); + let times = rows + .iter() + .map(|row| row["time"].as_str().unwrap()) + .collect::>(); + assert!( + times.is_sorted(), + "trace activity must be ordered by time: {times:?}" + ); rows } #[actix_web::test] async fn trace_activity_query_executes_with_distinct_bucketed_counts() { - let rows = run_trace_activity_fixture(&["service.name".to_owned()], &[]).await; - let summary = rows + let rows = run_trace_activity_fixture(&["service.name".to_owned()], &[], "BIGINT").await; + let mut summary = rows .iter() .map(|row| { ( @@ -2150,6 +2345,8 @@ mod tests { ) }) .collect::>(); + // Rows are ordered by bucket only; order groups within a bucket here. + summary.sort(); assert_eq!( summary, [ @@ -2175,12 +2372,18 @@ mod tests { &fields, ) .unwrap(); - let rows = run_trace_activity_fixture(&[], &clauses).await; - assert_eq!(rows.len(), 2); - assert!( - rows.iter() - .all(|row| row["span_count"] == 1 && row["trace_count"] == 1) - ); + // Status codes are integers in some trace datasets and Float64 in others. + for status_type in ["BIGINT", "DOUBLE"] { + let rows = run_trace_activity_fixture(&[], &clauses, status_type).await; + assert_eq!( + rows, + [ + json!({ "time": "2026-10-22T00:00:00", "span_count": 1, "trace_count": 1 }), + json!({ "time": "2026-10-22T00:01:00", "span_count": 1, "trace_count": 1 }), + ], + "{status_type}" + ); + } } #[actix_web::test] @@ -2219,6 +2422,94 @@ mod tests { assert!(bounded_tool_sql(&query).is_ok()); } + #[actix_web::test] + async fn trace_activity_requires_span_columns_and_uses_the_time_partition() { + let fields = + ["p_timestamp", "span_span_id", "span_trace_id", "event_time"].map(str::to_owned); + let time_field = |fields: &[String], partition: Option<&str>| { + super::trace_activity_time_field("traces", fields, partition.map(str::to_owned)) + }; + assert_eq!(time_field(&fields, None).unwrap(), "p_timestamp"); + assert_eq!( + time_field(&fields, Some("event_time")).unwrap(), + "event_time" + ); + assert_eq!( + time_field(&fields, Some("start_time")).unwrap_err(), + "trace dataset `traces` is missing time field `start_time`" + ); + for missing in ["span_span_id", "span_trace_id"] { + let fields = fields + .iter() + .filter(|field| *field != missing) + .cloned() + .collect::>(); + assert_eq!( + time_field(&fields, None).unwrap_err(), + format!("trace dataset `traces` is missing required field `{missing}`") + ); + } + } + + #[actix_web::test] + async fn dataset_activity_summary_orders_truncates_and_reports_failures() { + use super::DatasetActivity; + let count = |name: &str, events| DatasetActivity::Count { + name: name.to_owned(), + events, + }; + let error = |index: usize| DatasetActivity::Error { + name: format!("broken-{index}"), + message: "unavailable".to_owned(), + }; + let summarize = |activity, limit| { + super::summarize_dataset_activity( + activity, + limit, + 30, + "logs", + "2026-10-22T00:00:00Z", + "2026-10-22T01:00:00Z", + ) + }; + + let summary = summarize( + vec![ + count("b", 5), + count("a", 5), + count("c", 0), + count("d", 9), + DatasetActivity::Skipped, + error(0), + ], + 2, + ); + assert_eq!(summary["status"], "resolved"); + assert_eq!(summary["selectedDatasets"], json!(["d", "a"])); + assert_eq!( + summary["datasets"], + json!([ + { "name": "d", "eventsInRange": 9 }, + { "name": "a", "eventsInRange": 5 } + ]) + ); + assert_eq!(summary["activeDatasetCount"], 3); + assert_eq!(summary["truncated"], true); + assert_eq!(summary["queriedDatasets"], 4); + assert_eq!(summary["visibleDatasets"], 30); + assert_eq!(summary["unavailableDatasets"].as_array().unwrap().len(), 1); + + let quiet = summarize(vec![count("a", 0), DatasetActivity::Skipped], 5); + assert_eq!(quiet["status"], "not_found"); + assert_eq!(quiet["truncated"], false); + + let failed = summarize((0..25).map(error).chain([count("a", 0)]).collect(), 5); + assert_eq!(failed["status"], "error"); + assert!(failed.get("datasets").is_none()); + assert_eq!(failed["queriedDatasets"], 1); + assert_eq!(failed["unavailableDatasets"].as_array().unwrap().len(), 20); + } + #[actix_web::test] async fn trace_activity_bucket_interval_is_bounded() { assert_eq!( @@ -2260,6 +2551,56 @@ mod tests { ); } + #[actix_web::test] + async fn sample_events_selects_at_most_ten_columns() { + let fields = (0..15) + .map(|index| format!("f{index}")) + .chain(["p_timestamp".to_owned()]) + .collect::>(); + let query = + sample_events_query(&json!({ "dataset": "logs" }), &fields, Utc::now()).unwrap(); + let projection = query + .query + .strip_prefix("SELECT ") + .and_then(|query| query.split_once(" FROM ")) + .unwrap() + .0 + .split(", ") + .collect::>(); + assert_eq!( + projection, + [ + "\"p_timestamp\"", + "\"f0\"", + "\"f1\"", + "\"f2\"", + "\"f3\"", + "\"f4\"", + "\"f5\"", + "\"f6\"", + "\"f7\"", + "\"f8\"" + ] + ); + assert!( + query + .query + .ends_with(" ORDER BY \"p_timestamp\" DESC LIMIT 10") + ); + } + + #[actix_web::test] + async fn time_resolution_rejects_durations_beyond_the_supported_range() { + let reference = DateTime::parse_from_rfc3339("2026-10-01T12:00:00Z") + .unwrap() + .to_utc(); + for expression in ["last 1000000 years", "now-300000y"] { + let error = + resolve_time_range_at(&json!({ "expression": expression }), reference).unwrap_err(); + assert!(error.contains("too large"), "{expression}: {error}"); + } + } + #[actix_web::test] async fn sample_events_requires_both_explicit_range_bounds() { let fields = ["p_timestamp".to_owned()]; From 77723f9780b0fe15f9da2cbe8f4629afe6b76457 Mon Sep 17 00:00:00 2001 From: VibeSWE <338779468+VibeSWE@users.noreply.github.com> Date: Sat, 24 Oct 2026 14:25:04 -0500 Subject: [PATCH 5/5] docs: make tool API claims match the code - Listed inputSchemas omit root combinators; the server still enforces them. - SQL limits apply to the outermost projection and top-level set-operation branches, not CTEs or subqueries; a semicolon followed by a comment fails. - resolve_dataset: case-insensitive names, the `error` status, and scoring. - get_log_context: bound pairing and exclusivity, minute rounding of windows, and the anchor range check. - get_trace_activity: `query` is the inner SQL without a time predicate; rerun it through the query API with the same startTime and endTime. Co-Authored-By: Claude Opus 5.5 --- docs/tool-api.md | 83 +++++++++++++++++++++++++++++++++++------------- 1 file changed, 61 insertions(+), 22 deletions(-) diff --git a/docs/tool-api.md b/docs/tool-api.md index fc9cca74f..74d415fdd 100644 --- a/docs/tool-api.md +++ b/docs/tool-api.md @@ -36,6 +36,11 @@ Response: } ``` +A listed `inputSchema` never has a root-level `oneOf`, `anyOf` or `allOf`, +because several LLM tool APIs reject tool input schemas with root combinators. +Where the catalog schema has one, the server still enforces it when the tool is +called, and the affected property descriptions state the constraint. + The response includes an `ETag` and `Cache-Control: private, max-age=300, must-revalidate`. A request with a matching `If-None-Match` receives `304 Not Modified`. @@ -135,11 +140,15 @@ and pagination are not identical to this summary tool. `query_sql`, `explain_query` and other tools that run SQL internally parse the statement with DataFusion's SQL parser and accept exactly one `SELECT` or `WITH` -query. A trailing semicolon and trailing comments are allowed. Every `SELECT` -projection, including those inside nested queries and set operations, must -name at most 10 explicit columns. `*` and `table.*` are rejected. Casts to -`INT64`, `UINT64` or `FLOAT64` are rejected with a hint to use `BIGINT` or -`DOUBLE`. The statement is wrapped with an outer `LIMIT 500`. `sample_events` +query. A trailing semicolon or trailing comments are allowed, but a semicolon +followed by a comment passes validation and then fails when the query runs. +The outermost projection must name at most 10 explicit columns, and `*` and +`table.*` are rejected there. The same limit applies to each branch of a +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` returns at most 10 columns, choosing common observability fields such as `p_timestamp`, `service.name`, `severity_text`, `body` and `trace_id` before other fields. @@ -150,13 +159,25 @@ other fields. datasets. `resolve_dataset` accepts an optional `telemetryType` (`logs`, `metrics` or -`traces`) and `limit` (1–20, default 5). It resolves a `selectedDataset` or -`contextDataset` exactly, or resolves the only visible dataset. Otherwise it -ranks visible datasets by name match against `query`. It resolves only a -unique high-confidence name match. In every other case it returns -`status: "ambiguous"` with up to `limit` scored `candidates`, or `"not_found"`. -It never falls back to an arbitrary dataset. A selected dataset whose telemetry -type does not match `telemetryType` returns `not_found`. +`traces`) and `limit` (1–20, default 5). Its `status` is `resolved`, +`ambiguous`, `not_found` or `error`. Dataset names are matched +case-insensitively. + +- A `selectedDataset` or `contextDataset` resolves to the visible dataset with + that name. If no visible dataset has that name, the result has + `status: "error"`. If its telemetry type does not match `telemetryType`, the + result is `not_found`. +- Otherwise, if exactly one dataset is visible, it resolves to that dataset, or + returns `not_found` if its telemetry type does not match. +- Otherwise it scores the visible datasets that match `telemetryType` by name + against `query`: 1000 if the name equals the query, 750 if the query contains + the name, or 500 if the name contains the query, plus 25 for each query term + longer than two characters that the name contains, plus 15 when + `telemetryType` is given. It resolves only a unique top score of at least + 500. In every other case it returns `ambiguous` with up to `limit` scored + `candidates`, or `not_found` if no visible dataset matches `telemetryType`. + +It never falls back to an arbitrary dataset. `discover_datasets` counts events in the requested range for each visible dataset of the requested telemetry type that the caller can query. It returns @@ -177,12 +198,23 @@ provided together. They take precedence over `minutes`. `get_log_context` uses the same implementation as `POST /api/v1/query/context`. Arguments must include exactly one of `log`, `body` or `message`, holding the -anchor row's exact value. The schema enforces this with `oneOf`, so calls with -none or more than one are rejected with `400`. If no context bounds are given, -the window defaults to `contextWindow: "5 minutes"`. Returned records keep at -most 10 fields per row, preferring `p_timestamp`, service, severity, message -and trace identifiers. The response also includes the anchor index, scope and -the executed queries. +anchor row's exact value. The catalog schema enforces this with a root `oneOf`, +so calls with none or more than one are rejected with `400`. The listed +`inputSchema` omits that `oneOf` (see [List tools](#list-tools)) and states the +rule in the property descriptions. + +Context bounds are either `contextWindow` or both `contextStartTime` and +`contextEndTime`. Giving only one explicit bound, or a window together with an +explicit bound, is an error. If no bounds are given, the tool uses +`contextWindow: "5 minutes"`. A window extends that duration on each side of +`pTimestamp`, and both ends are rounded down to the minute; if they round to the +same minute, the end is one minute after the start. Explicit bounds are used at +millisecond precision without rounding, and `pTimestamp` must be at or after +the start and before the end. `pageSize` is 1–500, default 100. + +Returned records keep at most 10 fields per row, preferring `p_timestamp`, +service, severity, message and trace identifiers. The response also includes +the anchor index, scope and the executed queries. ## Trace activity @@ -211,9 +243,14 @@ 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. The result contains the exact executed SQL as `query`, -along with `bucketInterval`, `semantics` and up to 500 `rows`. Clients can -rerun that SQL unchanged, for example in a dashboard tile. +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`. ## Argument schema validation @@ -221,4 +258,6 @@ rerun that SQL unchanged, for example in a dashboard tile. `additionalProperties`, `enum`, `const`, `minLength`/`maxLength`, `minimum`/`maximum`, `minItems`/`maxItems`, `items`, `anyOf` and `oneOf`. `oneOf` passes only when exactly one variant matches. Combinators are checked -in addition to sibling keywords; they do not replace them. +in addition to sibling keywords; they do not replace them. Validation uses the +catalog schema, including any root combinators omitted from the listed +`inputSchema`.