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..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`. @@ -130,3 +135,129 @@ 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 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. + +## 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). 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 +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 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 + +`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. 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 + +`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 +in addition to sibling keywords; they do not replace them. Validation uses the +catalog schema, including any root combinators omitted from the listed +`inputSchema`. diff --git a/src/handlers/http/query_context.rs b/src/handlers/http/query_context.rs index 3b4229c29..23e768ddd 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 { @@ -364,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)) @@ -1190,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 1ec3625aa..cb7ecc710 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") @@ -184,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(), @@ -197,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 "))) } } @@ -357,7 +371,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 +423,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 +433,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 +557,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", @@ -811,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)) @@ -819,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; } @@ -839,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] @@ -848,7 +910,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 +949,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 +1033,192 @@ 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 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| { + 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() + ); + 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!( + 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 +1276,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] @@ -1054,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}), @@ -1067,6 +1372,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!({ diff --git a/src/tool_executor.rs b/src/tool_executor.rs index 772acac3d..360632335 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,110 @@ 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}"))?; + 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() { + 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 +335,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,79 +516,189 @@ 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; + + 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(); + 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 })); + } } } - Ok(json!({ + 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 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::>(); + 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::>(), + }) } fn schema_tool_result(schema: &arrow_schema::Schema, arguments: &Value) -> Result { @@ -318,110 +752,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 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)); + match_fields += 1; + } + } + if match_fields != 1 { + return Err( + "get_log_context requires exactly 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 +1329,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 +1351,7 @@ async fn execute_trace_tool( }, context.tenant_id, context.session_key, - Some(context.request_headers.clone()), + Some(context.query_headers.clone()), ) .await }; @@ -839,14 +1369,206 @@ async fn execute_trace_tool( Ok(json!({ "results": results, "errors": errors })) } -async fn execute_identity_tool( - name: &str, - arguments: &Value, - context: OssToolContext<'_>, -) -> Result { - match name { - "list_users" => { - authorize(context.session_key, Action::ListUser, None, None)?; +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::>(); + 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!( + "trace dataset `{dataset}` is missing required field `{required}`" + )); + } + } + 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(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, + context: OssToolContext<'_>, +) -> Result { + match name { + "list_users" => { + authorize(context.session_key, Action::ListUser, None, None)?; Ok(rbac::list_users_internal(context.tenant_id)) } "get_user" => { @@ -1009,12 +1731,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 +1769,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 +1796,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 +1815,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 +1925,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(); @@ -1168,9 +1945,12 @@ mod tests { use chrono::{DateTime, Duration, Utc}; use serde_json::{Value, json}; + use crate::handlers::http::query_context::LogContextRequest; + 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 +1990,628 @@ 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 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!([ + { "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() { + // 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] + 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_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" + })) + .unwrap(); + assert_eq!(request[anchor], "failed"); + } + } + + #[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(*)")); + } + + /// 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(&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: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 + .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 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()], &[], "BIGINT").await; + let mut 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::>(); + // Rows are ordered by bucket only; order groups within a bucket here. + summary.sort(); + 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(); + // 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] + 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_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!( + 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_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()]; + 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 +2874,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 =