fix(querier): let the MCP query_logs tool run count aggregations

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<Stage>, 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<Stage> {
+        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.

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]