coderabbit comments fix

Author: nikhilsinhaparseable

src/handlers/http/query_context.rs

@@ -138,14 +138,26 @@ 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),
         &tenant_id,
     )
     .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<String>,
+) -> Result<LogContextResponse, QueryError> {
+    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<u64>) -> Result<u64, QueryError> {

src/tool_catalog.rs

@@ -136,7 +136,7 @@ pub fn validate_tool_arguments(schema: &Value, arguments: &Value) -> Result<(),
 
 fn validate_schema_value(schema: &Value, value: &Value, path: &str) -> Result<(), String> {
     if let Some(result) = validate_any_of(schema, value, path) {
-        return result;
+        result?;
     }
 
     validate_enum(schema, value, path)?;
@@ -337,7 +337,9 @@ pub fn oss_read_tool_specs() -> Vec<ToolSpec> {
                 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"],
             ),
@@ -391,21 +393,29 @@ pub fn oss_read_tool_specs() -> Vec<ToolSpec> {
             "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["anyOf"] = json!([
+                    { "required": ["log"] },
+                    { "required": ["body"] },
+                    { "required": ["message"] }
+                ]);
+                schema
+            },
         ),
         ToolSpec::new(
             "list_log_errors",
@@ -1043,6 +1053,42 @@ mod tests {
                 .validate_arguments(&serde_json::json!({ "dataset": "logs" }))
                 .is_ok()
         );
+
+        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()
+        );
     }
 
     #[actix_web::test]