From 1b38aa5704c83dce70ba64fa34663b6995e1868db0330c0ef7727b1b40523df0 Mon Sep 17 00:00:00 2001
From: larssand <larssand@LPW0JW6L1>
Date: Tue, 30 Jun 2026 11:47:47 +0200
Subject: [PATCH] fix aggregate search

---
 src/fgai/graylog_aggregate.py   | 27 +++++++++++++++++++++++++++
 src/fgai/monitor.py             |  4 ++--
 tests/test_graylog_aggregate.py |  1 +
 3 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/src/fgai/graylog_aggregate.py b/src/fgai/graylog_aggregate.py
index 6bbd0b0..6057e9c 100644
--- a/src/fgai/graylog_aggregate.py
+++ b/src/fgai/graylog_aggregate.py
@@ -57,6 +57,20 @@ class GraylogAggregateSource:
     def fetch_count(self, *, range_seconds: int = 300) -> dict[str, object]:
         status = self.client.probe()
         variants: list[dict[str, object]] = [
+            {
+                "query": self.query,
+                "streams": [self.stream] if self.stream else [],
+                "range_seconds": max(1, int(range_seconds)),
+                "group_by": [],
+                "metrics": ["count()"],
+            },
+            {
+                "query": self.query,
+                "streams": [self.stream] if self.stream else [],
+                "range_seconds": max(1, int(range_seconds)),
+                "group_by": [],
+                "metrics": ["count"],
+            },
             {
                 "query": self.query,
                 "streams": [self.stream] if self.stream else [],
@@ -64,6 +78,13 @@ class GraylogAggregateSource:
                 "group_by": [],
                 "metrics": [{"function": "count"}],
             },
+            {
+                "query": self.query,
+                "streams": [self.stream] if self.stream else [],
+                "range_seconds": max(1, int(range_seconds)),
+                "groups": [],
+                "series": ["count()"],
+            },
             {
                 "query": self.query,
                 "streams": [self.stream] if self.stream else [],
@@ -71,6 +92,12 @@ class GraylogAggregateSource:
                 "groups": [],
                 "series": [{"function": "count"}],
             },
+            {
+                "query": self.query,
+                "streams": [self.stream] if self.stream else [],
+                "range_seconds": max(1, int(range_seconds)),
+                "limit": 1,
+            },
             {
                 "query": self.query,
                 "streams": [self.stream] if self.stream else [],
diff --git a/src/fgai/monitor.py b/src/fgai/monitor.py
index 3a04907..3983b36 100644
--- a/src/fgai/monitor.py
+++ b/src/fgai/monitor.py
@@ -93,7 +93,7 @@ def _stream_coverage(runtime_values: dict[str, object], stream_profiles: dict[st
             "latest_event_time": str(status.get("latest_event_time", "")),
             "truncated": bool(status.get("truncated")),
             "partial": bool(status.get("partial")),
-            "error": str(status.get("error", "")),
+            "error": str(status.get("error", "") or status.get("aggregate_error", "")),
             "health": "not_enabled" if not enabled else "partial_fetch" if status.get("partial") else "missing_profile" if not profile else "no_events" if int(status.get("events_fetched", 0) or 0) == 0 else "learning" if total_fields and ready_fields < total_fields else "ready" if total_fields else "profile_needs_fields",
         })
     return rows
@@ -171,7 +171,7 @@ def build_status(
                 if truncated_streams:
                     warnings.append(f"{len(truncated_streams)} stream(s) hit max_events_per_stream; high EPS means the analysis window is only partially sampled.")
                 mcp_status = {
-                    "status": "partial" if partial_streams else "connected",
+                    "status": "partial" if partial_streams or aggregate_errors else "connected",
                     "streams": stream_statuses,
                     "events_fetched": len(events),
                     "raw_events_fetched": len(events),
diff --git a/tests/test_graylog_aggregate.py b/tests/test_graylog_aggregate.py
index 1b98968..0961586 100644
--- a/tests/test_graylog_aggregate.py
+++ b/tests/test_graylog_aggregate.py
@@ -27,6 +27,7 @@ class GraylogAggregateTests(unittest.TestCase):
         self.assertEqual(status["aggregate_status"], "ok")
         self.assertEqual(status["aggregate_events"], 12345)
         self.assertEqual(client.arguments[0]["streams"], ["firewall"])
+        self.assertEqual(client.arguments[0]["metrics"], ["count()"])
 
     def test_tries_fallback_argument_shape_after_tool_error(self):
         client = _AggregateClient([