add mcp fix

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]:
+    def probe(self, *, force_new_session: bool = False) -> dict[str, object]:
-        self.session_id = None
+        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

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),