add timeout mcp poll

src/fgai/graylog_aggregate.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import json
+import time
 from collections.abc import Iterable
 
 from .graylog_mcp import GraylogMcpClient
@@ -69,7 +70,7 @@ class GraylogAggregateSource:
         self.stream = stream
         self.query = query or "*"
 
-    def fetch_count(self, *, range_seconds: int = 300, probe_status: dict[str, object] | None = None) -> dict[str, object]:
+    def fetch_count(self, *, range_seconds: int = 300, probe_status: dict[str, object] | None = None, deadline_monotonic: float | None = None) -> dict[str, object]:
         if probe_status is not None:
             status = dict(probe_status)
         else:
@@ -133,6 +134,9 @@ class GraylogAggregateSource:
         seen_variants: set[str] = set()
         errors: list[str] = []
         for arguments in variants:
+            if deadline_monotonic is not None and time.monotonic() >= deadline_monotonic:
+                errors.append("poll_budget_exceeded")
+                break
             variant_key = json.dumps(arguments, sort_keys=True)
             if variant_key in seen_variants:
                 continue

src/fgai/graylog_source.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import json
+import time
 from collections.abc import Iterable
 
 from .graylog_mcp import GraylogMcpClient
@@ -68,7 +69,7 @@ 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, probe_status: dict[str, object] | None = None) -> 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, deadline_monotonic: float | None = None) -> tuple[list[LogEvent], dict[str, object]]:
         if probe_status is not None:
             status = dict(probe_status)
         else:
@@ -98,6 +99,9 @@ class GraylogStreamSource:
         pages = 0
         partial_error = ""
         while len(events) < max_events:
+            if deadline_monotonic is not None and time.monotonic() >= deadline_monotonic:
+                partial_error = "poll_budget_exceeded"
+                break
             request_limit = min(page_size, max_events - len(events))
             try:
                 result = self.client.call_tool("search_messages", {**arguments, "limit": request_limit, "offset": len(events)})

src/fgai/monitor.py

@@ -176,11 +176,11 @@ def build_status(
                 probe_status = client.probe()
                 discovery_store = FieldDiscoveryStore(history_path) if history_path else None
                 catalog_fields_total = 0
-                for stream_config in stream_configs:
+                def budget_exceeded() -> bool:
-                    stream_id = str(stream_config["id"])
+                    return time.monotonic() >= poll_deadline
-                    stream_name = str(stream_config.get("title", "") or stream_titles.get(stream_id) or stream_id)
+
-                    if time.monotonic() >= poll_deadline:
+                def skipped_status(stream_id: str, stream_name: str) -> dict[str, object]:
-                        stream_statuses.append({
+                    return {
                         "stream_id": stream_id,
                         "stream_name": stream_name,
                         "source": "graylog_mcp",
@@ -193,7 +193,13 @@ def build_status(
                         "latest_event_time": "",
                         "raw_sample_limit": min(raw_sample_events, 10_000) if use_aggregate else max_events_per_stream,
                         "health_detail": "Skipped because the MCP poll time budget was reached before this stream.",
-                        })
+                    }
+
+                for stream_config in stream_configs:
+                    stream_id = str(stream_config["id"])
+                    stream_name = str(stream_config.get("title", "") or stream_titles.get(stream_id) or stream_id)
+                    if budget_exceeded():
+                        stream_statuses.append(skipped_status(stream_id, stream_name))
                         continue
                     profile = stream_profiles.get(stream_id)
                     profile_fields = (
@@ -205,17 +211,20 @@ def build_status(
                         *tuple(str(getattr(relation, "left", "")) for relation in getattr(profile, "relationship_fields", ())),
                         *tuple(str(getattr(relation, "right", "")) for relation in getattr(profile, "relationship_fields", ())),
                     ) if profile else ()
-                    if discovery_store:
+                    if discovery_store and not budget_exceeded():
                         try:
                             catalog_fields_total += discovery_store.ingest_catalog(stream_id, stream_name, _graylog_fields_from_result(client.call_tool("list_fields", {"streams": [stream_id]})))
                         except RuntimeError:
                             pass
                     aggregate_status: dict[str, object] = {}
-                    if use_aggregate:
+                    if use_aggregate and not budget_exceeded():
-                        aggregate_status = GraylogAggregateSource(client, stream_id, str(runtime_values.get("graylog_query", "*"))).fetch_count(range_seconds=range_seconds, probe_status=probe_status)
+                        aggregate_status = GraylogAggregateSource(client, stream_id, str(runtime_values.get("graylog_query", "*"))).fetch_count(range_seconds=range_seconds, probe_status=probe_status, deadline_monotonic=poll_deadline)
                         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(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)
+                    if budget_exceeded():
+                        stream_statuses.append({**skipped_status(stream_id, stream_name), **aggregate_status})
+                        continue
+                    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, deadline_monotonic=poll_deadline)
                     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

@@ -1,4 +1,5 @@
 import unittest
+import time
 
 from fgai.graylog_aggregate import GraylogAggregateSource
 
@@ -92,6 +93,17 @@ class GraylogAggregateTests(unittest.TestCase):
         self.assertIn("groupings", client.arguments[0])
         self.assertNotIn("group_by", client.arguments[0])
 
+    def test_deadline_stops_aggregate_attempts(self):
+        client = _AggregateClient([
+            {"result": {"content": [{"type": "text", "text": '{"events": 9}'}]}}
+        ])
+
+        status = GraylogAggregateSource(client, "firewall").fetch_count(deadline_monotonic=time.monotonic() - 1)
+
+        self.assertEqual(status["aggregate_status"], "error")
+        self.assertIn("poll_budget_exceeded", status["aggregate_error"])
+        self.assertEqual(client.arguments, [])
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_graylog_source.py

@@ -1,4 +1,5 @@
 import unittest
+import time
 
 from fgai.graylog_source import GraylogStreamSource
 
@@ -73,6 +74,14 @@ class GraylogSourceTests(unittest.TestCase):
         self.assertTrue(status["partial"])
         self.assertIn("graylog_search_error", status["error"])
 
+    def test_deadline_stops_raw_fetch_before_search(self):
+        client = _Client()
+        events, status = GraylogStreamSource(client, "vpn").fetch(deadline_monotonic=time.monotonic() - 1)
+        self.assertEqual(events, [])
+        self.assertTrue(status["partial"])
+        self.assertEqual(status["error"], "poll_budget_exceeded")
+        self.assertIsNone(client.arguments)
+
     def test_returns_partial_status_instead_of_raising_on_probe_error(self):
         events, status = GraylogStreamSource(_ProbeErrorClient(), "vpn").fetch()
         self.assertEqual(events, [])