diff --git a/crates/ourios-server/src/mcp.rs b/crates/ourios-server/src/mcp.rs index 89b76425..8782cfbf 100644 --- a/crates/ourios-server/src/mcp.rs +++ b/crates/ourios-server/src/mcp.rs @@ -27,6 +27,7 @@ use ourios_core::tenant::TenantId; use ourios_ingester::receiver::{AuthBinding, AuthResolver}; use ourios_parquet::PromotedAttributes; use ourios_querier::Querier; +use ourios_querier::dsl::ir::Stage; use ourios_querier::dsl::{self, Statement}; use rmcp::handler::server::ServerHandler; use rmcp::handler::server::wrapper::Parameters; @@ -214,6 +215,18 @@ fn normalize_tenant(raw: &str) -> Result<&str, ErrorData> { Ok(tenant) } +/// Apply the tool's row `cap` to `stages`, unless the statement is a `count +/// [by …]` aggregation. `count` and `limit` are mutually exclusive +/// (`compile::validate`) — an aggregation answers with its grouped-count map, +/// not a capped row set, so injecting the cap would reject the query. Mirrors +/// the JSON API's guard (`querier::handle_query`). +fn cap_rows_unless_aggregation(stages: &mut Vec, cap: u64) { + let is_aggregation = stages.iter().any(|s| matches!(s, Stage::Count { .. })); + if !is_aggregation { + apply_limit(stages, cap, cap); + } +} + impl OuriosMcp { /// RFC 0026 per-call tenant binding: re-resolve the request's bearer /// (rmcp forwards the HTTP parts into the tool context) and require @@ -296,7 +309,7 @@ impl OuriosMcp { // `limit` stage inside the statement clamps to it, so the // documented "maximum rendered rows" contract holds. let cap = args.limit.unwrap_or(DEFAULT_LIMIT).min(MAX_LIMIT); - apply_limit(&mut query.stages, cap, cap); + cap_rows_unless_aggregation(&mut query.stages, cap); let tenant = TenantId::new(tenant_arg); let started = std::time::Instant::now(); let result = self @@ -625,7 +638,45 @@ mod tests { use ourios_querier::dsl::ir::SeverityName; use parquet::file::reader::{FileReader, SerializedFileReader}; - use super::{GRAMMAR_SECTION, query_schema_document}; + use super::{GRAMMAR_SECTION, Stage, query_schema_document}; + + /// Parse a logs DSL statement to its stage list (guard-test helper). + fn logs_stages(query: &str) -> Vec { + match super::dsl::parse_statement(query).expect("valid dsl") { + super::Statement::Logs(q) => q.stages, + super::Statement::Drift(_) => panic!("expected a logs query, not a drift statement"), + } + } + + /// `cap_rows_unless_aggregation`: a row query gets the tool's cap + /// injected as a `limit` stage — the "maximum rendered rows" contract. + #[test] + fn row_query_gets_the_row_cap() { + let mut stages = logs_stages("template_id == 1"); + super::cap_rows_unless_aggregation(&mut stages, 10); + assert!( + stages.iter().any(|s| matches!(s, Stage::Limit(10))), + "a non-aggregation gets the cap injected: {stages:?}", + ); + } + + /// A `count [by …]` aggregation is left uncapped — `count` and `limit` + /// are mutually exclusive, so injecting the cap would reject the query + /// (the bug this guard fixes). Covers both bare `count` and `count by`. + #[test] + fn count_aggregation_is_left_uncapped() { + for query in [ + "template_id == 1 | count", + "template_id == 1 | count by template_id", + ] { + let mut stages = logs_stages(query); + super::cap_rows_unless_aggregation(&mut stages, 10); + assert!( + !stages.iter().any(|s| matches!(s, Stage::Limit(_))), + "an aggregation keeps no limit stage ({query}): {stages:?}", + ); + } + } /// The extraction invariants RFC0027.6 leans on: heading-first, /// non-empty, and bounded before the next top-level section. diff --git a/crates/ourios-server/tests/it/rfc0027_mcp.rs b/crates/ourios-server/tests/it/rfc0027_mcp.rs index 92c1a8ef..5e0c9f9d 100644 --- a/crates/ourios-server/tests/it/rfc0027_mcp.rs +++ b/crates/ourios-server/tests/it/rfc0027_mcp.rs @@ -362,6 +362,56 @@ async fn rfc0027_3_query_logs() { ); } +/// Scenario RFC0027.3 — `query_logs`: equivalence with the JSON API. +/// +/// Regression facet: a `count [by …]` aggregation. The tool auto-attaches +/// its row cap, but `count` and `limit` are mutually exclusive +/// (`compile::validate`), so the cap must be skipped for aggregations — +/// mirroring the JSON API's `is_aggregation` guard. Before the fix the tool +/// rejected every `count by` with "does not support `limit`", so no +/// aggregation could run through the MCP surface. +/// See `docs/rfcs/0027-mcp-query-surface.md` §5. +#[tokio::test] +async fn rfc0027_3_query_logs_count_aggregation() { + let bucket = tempfile::tempdir().expect("temp"); + crate::rfc0016_query_endpoint::seed_two_records(bucket.path()); + + let router = ourios_server::querier::router_with_mcp( + bucket.path().to_path_buf(), + crate::rfc0016_query_endpoint::SHARED_HUGE_WINDOW, + ourios_ingester::receiver::AuthResolver::static_only(None), + true, + ); + let query = "template_id == 1 | count by template_id"; + // A non-null `limit` is passed on purpose: the fix must ignore it for an + // aggregation rather than inject it and reject the query. + let body = mcp_tool_call( + &router, + None, + "query_logs", + serde_json::json!({"tenant": "acme", "query": query, "limit": 10}), + ) + .await; + let mcp_payload = tool_text(&body); + assert!( + mcp_payload["aggregate"] + .as_array() + .is_some_and(|groups| !groups.is_empty()), + "the count aggregation returns a non-empty grouped-count map, not a tool error: {body}", + ); + + // Identical to the JSON API for the same statement — the adapter adds + // nothing but the protocol, aggregations included. + let (status, json_payload) = + crate::rfc0016_query_endpoint::post_for_equivalence(bucket.path(), Some("acme"), query) + .await; + assert_eq!(status, StatusCode::OK); + assert_eq!( + mcp_payload, json_payload, + "MCP and JSON API agree on the aggregation", + ); +} + /// Scenario RFC0027.4 — `list_templates` matches the RFC 0017 registry. /// See `docs/rfcs/0027-mcp-query-surface.md` §5. #[tokio::test]