diff --git a/src/fgai/graylog_aggregate.py b/src/fgai/graylog_aggregate.py index 2119605..d6bacc7 100644 --- a/src/fgai/graylog_aggregate.py +++ b/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]: - try: - status = self.client.probe() - except RuntimeError as exc: - return { - "status": "error", - "source": "graylog_mcp_aggregate", - "aggregate_status": "error", - "aggregate_events": 0, - "aggregate_records": 0, - "aggregate_error": f"probe_error: {exc}", - "aggregate_errors": [f"probe_error: {exc}"], - "aggregate_schema_properties": [], - } + def fetch_count(self, *, range_seconds: int = 300, probe_status: dict[str, object] | None = None) -> 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_aggregate", + "aggregate_status": "error", + "aggregate_events": 0, + "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) diff --git a/src/fgai/graylog_source.py b/src/fgai/graylog_source.py index 35a5d6b..4a9dbaf 100644 --- a/src/fgai/graylog_source.py +++ b/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, diff --git a/src/fgai/monitor.py b/src/fgai/monitor.py index 5b3970a..004ab11 100644 --- a/src/fgai/monitor.py +++ b/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] diff --git a/tests/test_graylog_aggregate.py b/tests/test_graylog_aggregate.py index e1acef0..57f0fcc 100644 --- a/tests/test_graylog_aggregate.py +++ b/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"}]}}, diff --git a/tests/test_graylog_source.py b/tests/test_graylog_source.py index 8f5b06a..6edca5f 100644 --- a/tests/test_graylog_source.py +++ b/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)