Public/fgAI
SHA256
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add timeout mcp poll

Browse Source
This commit is contained in:
larssand
2026-07-02 21:17:50 +02:00
parent f45c275e20
commit 97d017e43c
5 changed files with 59 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
src/fgai/graylog_aggregate.py
View File

@@ -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

6
src/fgai/graylog_source.py
View File

@@ -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)})

29
src/fgai/monitor.py
View File

@@ -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:
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({
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",
@@ -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:
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]

12
tests/test_graylog_aggregate.py
View File

@@ -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()

9
tests/test_graylog_source.py
View File

@@ -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, [])