From 97d017e43c7f0f12ed87ce6e4de2d14b7efef7f74402f38da0b2b872633198af Mon Sep 17 00:00:00 2001 From: larssand Date: Thu, 2 Jul 2026 21:17:50 +0200 Subject: [PATCH] add timeout mcp poll --- src/fgai/graylog_aggregate.py | 6 ++++- src/fgai/graylog_source.py | 6 ++++- src/fgai/monitor.py | 47 ++++++++++++++++++++------------- tests/test_graylog_aggregate.py | 12 +++++++++ tests/test_graylog_source.py | 9 +++++++ 5 files changed, 59 insertions(+), 21 deletions(-) diff --git a/src/fgai/graylog_aggregate.py b/src/fgai/graylog_aggregate.py index d6bacc7..fd5044e 100644 --- a/src/fgai/graylog_aggregate.py +++ b/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 diff --git a/src/fgai/graylog_source.py b/src/fgai/graylog_source.py index ba1602a..165e565 100644 --- a/src/fgai/graylog_source.py +++ b/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)}) diff --git a/src/fgai/monitor.py b/src/fgai/monitor.py index 3aee78d..6f478a2 100644 --- a/src/fgai/monitor.py +++ b/src/fgai/monitor.py @@ -176,24 +176,30 @@ def build_status( probe_status = client.probe() discovery_store = FieldDiscoveryStore(history_path) if history_path else None catalog_fields_total = 0 + def budget_exceeded() -> bool: + return time.monotonic() >= poll_deadline + + def skipped_status(stream_id: str, stream_name: str) -> dict[str, object]: + return { + "stream_id": stream_id, + "stream_name": stream_name, + "source": "graylog_mcp", + "events_fetched": 0, + "aggregate_events": 0, + "pages": 0, + "partial": True, + "error": "skipped_poll_budget", + "truncated": False, + "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 time.monotonic() >= poll_deadline: - stream_statuses.append({ - "stream_id": stream_id, - "stream_name": stream_name, - "source": "graylog_mcp", - "events_fetched": 0, - "aggregate_events": 0, - "pages": 0, - "partial": True, - "error": "skipped_poll_budget", - "truncated": False, - "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.", - }) + 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: - aggregate_status = GraylogAggregateSource(client, stream_id, str(runtime_values.get("graylog_query", "*"))).fetch_count(range_seconds=range_seconds, probe_status=probe_status) + 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, 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] diff --git a/tests/test_graylog_aggregate.py b/tests/test_graylog_aggregate.py index 57f0fcc..758d4fb 100644 --- a/tests/test_graylog_aggregate.py +++ b/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() diff --git a/tests/test_graylog_source.py b/tests/test_graylog_source.py index 1bbf57d..a11ff15 100644 --- a/tests/test_graylog_source.py +++ b/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, [])