From 35c21332f3cc0be443d8c7fa8a94069aa86911d0985768220f476614c8531704 Mon Sep 17 00:00:00 2001 From: larssand Date: Sun, 21 Jun 2026 22:18:58 +0200 Subject: [PATCH] fix mcp search --- src/fgai/graylog_source.py | 11 ++++++++++- tests/test_graylog_source.py | 11 +++++++++-- 2 files changed, 19 insertions(+), 3 deletions(-) diff --git a/src/fgai/graylog_source.py b/src/fgai/graylog_source.py index c59b32e..b1a2ef0 100644 --- a/src/fgai/graylog_source.py +++ b/src/fgai/graylog_source.py @@ -44,8 +44,17 @@ class GraylogStreamSource: def fetch(self) -> tuple[list[LogEvent], dict[str, object]]: status = self.client.probe() - result = self.client.call_tool("search_messages", {"query": self.query, "stream": self.stream, "limit": 1000}) + arguments: dict[str, object] = {"query": self.query, "limit": 1000, "range_seconds": 300} + if self.stream: + arguments["streams"] = [self.stream] + result = self.client.call_tool("search_messages", arguments) content = result.get("result", {}).get("content", []) if isinstance(result.get("result"), dict) else [] + if isinstance(result.get("result"), dict) and result["result"].get("isError"): + detail = next( + (str(item.get("text")) for item in content if isinstance(item, dict) and item.get("type") == "text"), + "Graylog search failed", + ) + raise RuntimeError(f"graylog_search_error: {detail}") records: list[dict[str, object]] = [] for item in content if isinstance(content, list) else []: if isinstance(item, dict) and item.get("type") == "text": diff --git a/tests/test_graylog_source.py b/tests/test_graylog_source.py index 3138fe0..a5b5056 100644 --- a/tests/test_graylog_source.py +++ b/tests/test_graylog_source.py @@ -4,22 +4,29 @@ from fgai.graylog_source import GraylogStreamSource class _Client: + def __init__(self): + self.arguments = None + def probe(self): return {"status": "connected"} - def call_tool(self, _name, _arguments): + def call_tool(self, _name, arguments): + self.arguments = arguments return {"result": {"content": [{"type": "text", "text": '{"messages":[{"client":"10.0.0.1","server":"8.8.8.8","result":"deny"}]}' }]}} class GraylogSourceTests(unittest.TestCase): def test_applies_custom_mapping_to_generic_stream_message(self): + client = _Client() events, status = GraylogStreamSource( - _Client(), "vpn", field_mapping='{"srcip":"client","dstip":"server","action":"result"}' + client, "vpn", field_mapping='{"srcip":"client","dstip":"server","action":"result"}' ).fetch() self.assertEqual(status["events_fetched"], 1) self.assertEqual(events[0].src_ip, "10.0.0.1") self.assertEqual(events[0].dst_ip, "8.8.8.8") self.assertEqual(events[0].action, "deny") + self.assertEqual(client.arguments["streams"], ["vpn"]) + self.assertEqual(client.arguments["range_seconds"], 300) if __name__ == "__main__":