mcp issues

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]]:
-        try:
-            status = self.client.probe()
-        except RuntimeError as exc:
-            return [], {
-                "status": "error",
-                "source": "graylog_mcp",
-                "events_fetched": 0,
-                "pages": 0,
-                "partial": True,
-                "error": f"probe_error: {exc}",
-                "truncated": False,
-                "latest_event_time": "",
-            }
+    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]]:
+        if probe_status is not None:
+            status = dict(probe_status)
+        else:
+            try:
+                status = self.client.probe()
+            except RuntimeError as exc:
+                return [], {
+                    "status": "error",
+                    "source": "graylog_mcp",
+                    "events_fetched": 0,
+                    "pages": 0,
+                    "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,