mcp issues

src/fgai/graylog_aggregate.py

@@ -69,20 +69,23 @@ class GraylogAggregateSource:
         self.stream = stream
         self.query = query or "*"
 
-    def fetch_count(self, *, range_seconds: int = 300) -> dict[str, object]:
+    def fetch_count(self, *, range_seconds: int = 300, probe_status: dict[str, object] | None = None) -> dict[str, object]:
-        try:
+        if probe_status is not None:
-            status = self.client.probe()
+            status = dict(probe_status)
-        except RuntimeError as exc:
+        else:
-            return {
+            try:
-                "status": "error",
+                status = self.client.probe()
-                "source": "graylog_mcp_aggregate",
+            except RuntimeError as exc:
-                "aggregate_status": "error",
+                return {
-                "aggregate_events": 0,
+                    "status": "error",
-                "aggregate_records": 0,
+                    "source": "graylog_mcp_aggregate",
-                "aggregate_error": f"probe_error: {exc}",
+                    "aggregate_status": "error",
-                "aggregate_errors": [f"probe_error: {exc}"],
+                    "aggregate_events": 0,
-                "aggregate_schema_properties": [],
+                    "aggregate_records": 0,
-            }
+                    "aggregate_error": f"probe_error: {exc}",
+                    "aggregate_errors": [f"probe_error: {exc}"],
+                    "aggregate_schema_properties": [],
+                }
         tool_schemas = status.get("tool_schemas", {}) if isinstance(status.get("tool_schemas"), dict) else {}
         aggregate_schema = tool_schemas.get("aggregate_messages", {}) if isinstance(tool_schemas, dict) else {}
         properties = _schema_properties(aggregate_schema)

src/fgai/graylog_source.py

@@ -45,20 +45,23 @@ class GraylogStreamSource:
         if not isinstance(self.mapping, dict):
             raise RuntimeError("invalid_graylog_field_mapping")
 
-    def fetch(self, *, max_events: int = 5_000, range_seconds: int = 300) -> tuple[list[LogEvent], dict[str, object]]:
+    def fetch(self, *, max_events: int = 5_000, range_seconds: int = 300, probe_status: dict[str, object] | None = None) -> tuple[list[LogEvent], dict[str, object]]:
-        try:
+        if probe_status is not None:
-            status = self.client.probe()
+            status = dict(probe_status)
-        except RuntimeError as exc:
+        else:
-            return [], {
+            try:
-                "status": "error",
+                status = self.client.probe()
-                "source": "graylog_mcp",
+            except RuntimeError as exc:
-                "events_fetched": 0,
+                return [], {
-                "pages": 0,
+                    "status": "error",
-                "partial": True,
+                    "source": "graylog_mcp",
-                "error": f"probe_error: {exc}",
+                    "events_fetched": 0,
-                "truncated": False,
+                    "pages": 0,
-                "latest_event_time": "",
+                    "partial": True,
-            }
+                    "error": f"probe_error: {exc}",
+                    "truncated": False,
+                    "latest_event_time": "",
+                }
         mapping_fields = [str(value) for value in self.mapping.values() if isinstance(value, str)]
         arguments: dict[str, object] = {
             "query": self.query,

src/fgai/monitor.py

@@ -144,6 +144,8 @@ def build_status(
                 fetch_mode = str(runtime_values.get("graylog_fetch_mode", "auto") or "auto")
                 use_aggregate = fetch_mode == "aggregate" or (fetch_mode == "auto" and max_events_per_stream > raw_sample_events)
                 aggregate_events_total = 0
+                client = GraylogMcpClient(url, token, verify_tls=verify_tls)
+                probe_status = client.probe()
                 for stream_config in stream_configs:
                     stream_id = str(stream_config["id"])
                     profile = stream_profiles.get(stream_id)
@@ -157,10 +159,10 @@ def build_status(
                     stream_name = str(stream_config.get("title", "") or stream_titles.get(stream_id) or stream_id)
                     aggregate_status: dict[str, object] = {}
                     if use_aggregate:
-                        aggregate_status = GraylogAggregateSource(GraylogMcpClient(url, token, verify_tls=verify_tls), stream_id, str(runtime_values.get("graylog_query", "*"))).fetch_count(range_seconds=range_seconds)
+                        aggregate_status = GraylogAggregateSource(client, stream_id, str(runtime_values.get("graylog_query", "*"))).fetch_count(range_seconds=range_seconds, probe_status=probe_status)
                         aggregate_events_total += int(aggregate_status.get("aggregate_events", 0) or 0)
                     raw_limit = min(raw_sample_events, 10_000) if use_aggregate else max_events_per_stream
-                    stream_events, stream_status = GraylogStreamSource(GraylogMcpClient(url, token, verify_tls=verify_tls), stream_id, str(runtime_values.get("graylog_query", "*")), str(runtime_values.get("graylog_field_mapping", "")), stream_name, profile_fields).fetch(max_events=raw_limit, range_seconds=range_seconds)
+                    stream_events, stream_status = GraylogStreamSource(client, stream_id, str(runtime_values.get("graylog_query", "*")), str(runtime_values.get("graylog_field_mapping", "")), stream_name, profile_fields).fetch(max_events=raw_limit, range_seconds=range_seconds, probe_status=probe_status)
                     events.extend(stream_events)
                     stream_statuses.append({"stream_id": stream_id, "stream_name": stream_name, **aggregate_status, **stream_status, "raw_sample_limit": raw_limit})
                 sample_limited_streams = [item for item in stream_statuses if item.get("truncated") and use_aggregate]

tests/test_graylog_aggregate.py

@@ -8,8 +8,10 @@ class _AggregateClient:
         self.responses = list(responses)
         self.arguments = []
         self.schema = schema
+        self.probes = 0
 
     def probe(self):
+        self.probes += 1
         status = {"status": "connected"}
         if self.schema:
             status["tool_schemas"] = {"aggregate_messages": self.schema}
@@ -43,6 +45,14 @@ class GraylogAggregateTests(unittest.TestCase):
         self.assertEqual(status["aggregate_status"], "error")
         self.assertIn("probe_error", status["aggregate_error"])
 
+    def test_uses_existing_probe_status_without_reprobing(self):
+        client = _AggregateClient([
+            {"result": {"content": [{"type": "text", "text": '{"events": 9}'}]}}
+        ])
+        status = GraylogAggregateSource(client, "firewall").fetch_count(probe_status={"status": "connected"})
+        self.assertEqual(status["aggregate_events"], 9)
+        self.assertEqual(client.probes, 0)
+
     def test_tries_fallback_argument_shape_after_tool_error(self):
         client = _AggregateClient([
             {"result": {"isError": True, "content": [{"type": "text", "text": "bad metrics"}]}},

tests/test_graylog_source.py

@@ -6,8 +6,10 @@ from fgai.graylog_source import GraylogStreamSource
 class _Client:
     def __init__(self):
         self.arguments = None
+        self.probes = 0
 
     def probe(self):
+        self.probes += 1
         return {"status": "connected"}
 
     def call_tool(self, _name, arguments):
@@ -48,6 +50,11 @@ class GraylogSourceTests(unittest.TestCase):
         GraylogStreamSource(client, "vpn").fetch(range_seconds=86_400)
         self.assertEqual(client.arguments["range_seconds"], 86_400)
 
+    def test_uses_existing_probe_status_without_reprobing(self):
+        client = _Client()
+        GraylogStreamSource(client, "vpn").fetch(probe_status={"status": "connected"})
+        self.assertEqual(client.probes, 0)
+
     def test_marks_stream_truncated_when_max_events_reached(self):
         client = _Client()
         events, status = GraylogStreamSource(client, "vpn").fetch(max_events=2)