From 6a63de0e373198313b85d231cd3b4ea794f6597d8008cb4d0517b44e5dd9e5c0 Mon Sep 17 00:00:00 2001
From: larssand <larssand@LPW0JW6L1>
Date: Tue, 30 Jun 2026 21:30:30 +0200
Subject: [PATCH] add mcp fix

---
 src/fgai/graylog_mcp.py   |  7 ++++---
 tests/test_graylog_mcp.py | 14 ++++++++++++++
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/src/fgai/graylog_mcp.py b/src/fgai/graylog_mcp.py
index 16fee1c..5b61423 100644
--- a/src/fgai/graylog_mcp.py
+++ b/src/fgai/graylog_mcp.py
@@ -66,8 +66,9 @@ class GraylogMcpClient:
         credentials = base64.b64encode(f"{value}:token".encode("utf-8")).decode("ascii")
         return f"Basic {credentials}"
 
-    def probe(self) -> dict[str, object]:
-        self.session_id = None
+    def probe(self, *, force_new_session: bool = False) -> dict[str, object]:
+        if force_new_session:
+            self.session_id = None
         initialized = self._call(
             "initialize",
             {"protocolVersion": "2025-06-18", "capabilities": {}, "clientInfo": {"name": "fgAI", "version": "0.1"}},
@@ -90,7 +91,7 @@ class GraylogMcpClient:
         except RuntimeError as exc:
             if not self._should_reconnect(str(exc)):
                 raise
-            self.probe()
+            self.probe(force_new_session=True)
             return self._call("tools/call", {"name": name, "arguments": arguments})
 
     @staticmethod
diff --git a/tests/test_graylog_mcp.py b/tests/test_graylog_mcp.py
index c04b60d..d57c1d5 100644
--- a/tests/test_graylog_mcp.py
+++ b/tests/test_graylog_mcp.py
@@ -39,6 +39,20 @@ class GraylogMcpTests(unittest.TestCase):
         self.assertEqual(status["status"], "connected")
         self.assertIn("search_messages", status["tools"])
 
+    def test_probe_keeps_existing_session_unless_forced(self):
+        responses = iter([
+            _Response({"result": {"serverInfo": {"version": "7.1.6"}}}, session="session-1"),
+            _Response({}, session="session-1"),
+            _Response({"result": {"tools": []}}, session="session-1"),
+        ])
+        client = GraylogMcpClient("http://graylog/api/mcp", "raw-token")
+        client.session_id = "existing-session"
+
+        with patch("fgai.graylog_mcp.request.urlopen", side_effect=responses):
+            client.probe()
+
+        self.assertEqual(client.session_id, "session-1")
+
     def test_call_tool_reinitializes_after_stale_session_error(self):
         responses = iter([
             HTTPError("http://graylog/api/mcp", 404, "stale session", {}, None),