From 2287d04c4d2115f5ba9c371219539ec0722897f5 Mon Sep 17 00:00:00 2001
From: Heet Ranpura <heetranpura1@gmail.com>
Date: Fri, 31 Jul 2026 23:39:59 +0530
Subject: [PATCH] zulip: Reset session on ConnectionError to fix stale
 connections.

When a network middlebox silently drops an idle TCP connection, the next
request raises a ConnectionError. Previously the retry logic would reuse
the same dead session, failing all retry attempts because the connection
was never refreshed.

Now we close and null out the session on ConnectionError so
ensure_session() creates a fresh one on the next retry.

Fixes #761.
---
 zulip/tests/test_do_api_query.py | 83 ++++++++++++++++++++++++++++++++
 zulip/zulip/__init__.py          | 14 ++++--
 2 files changed, 94 insertions(+), 3 deletions(-)
 create mode 100644 zulip/tests/test_do_api_query.py

diff --git a/zulip/tests/test_do_api_query.py b/zulip/tests/test_do_api_query.py
new file mode 100644
index 000000000..56a0758b3
--- /dev/null
+++ b/zulip/tests/test_do_api_query.py
@@ -0,0 +1,83 @@
+from unittest import TestCase
+from unittest.mock import MagicMock, patch
+
+import requests
+
+import zulip
+
+
+class TestDoApiQuery(TestCase):
+    @patch("zulip.Client.get_server_settings")
+    def _make_client(self, mock_server_settings: MagicMock) -> zulip.Client:
+        mock_server_settings.return_value = {
+            "zulip_version": "1.0",
+            "zulip_feature_level": 0,
+        }
+        client = zulip.Client(
+            email="test@example.com",
+            api_key="fake-key",
+            site="https://example.com",
+        )
+        client.has_connected = True
+        return client
+
+    def test_session_reset_on_connection_error(self) -> None:
+        """When a ConnectionError occurs and the client has previously
+        connected, the stale session should be closed and replaced with
+        a fresh one so that the retry succeeds.
+        """
+        client = self._make_client()
+
+        stale_session = MagicMock(spec=requests.Session)
+        stale_session.request.side_effect = requests.exceptions.ConnectionError(
+            "Connection reset by peer"
+        )
+
+        fresh_session = MagicMock(spec=requests.Session)
+        mock_response = MagicMock()
+        mock_response.status_code = 200
+        mock_response.json.return_value = {"result": "success", "msg": ""}
+        fresh_session.request.return_value = mock_response
+
+        # Track how many sessions ensure_session has built.
+        sessions_created = 0
+
+        def mock_ensure_session(self: zulip.Client) -> None:
+            nonlocal sessions_created
+            if self.session is None:
+                sessions_created += 1
+                if sessions_created == 1:
+                    self.session = stale_session
+                else:
+                    self.session = fresh_session
+
+        with patch.object(zulip.Client, "ensure_session", mock_ensure_session):
+            result = client.do_api_query({}, "v1/messages", method="GET")
+
+        # The stale session should have been closed.
+        stale_session.close.assert_called_once()
+
+        # A fresh session should have been created for the retry.
+        self.assertEqual(sessions_created, 2)
+
+        # The request should ultimately succeed.
+        self.assertEqual(result["result"], "success")
+
+    def test_session_not_reset_when_never_connected(self) -> None:
+        """When the client has never successfully connected, a
+        ConnectionError should raise UnrecoverableNetworkError
+        without resetting the session.
+        """
+        client = self._make_client()
+        client.has_connected = False
+
+        mock_session = MagicMock(spec=requests.Session)
+        mock_session.request.side_effect = requests.exceptions.ConnectionError("Connection refused")
+        client.session = mock_session
+
+        with self.assertRaises(zulip.UnrecoverableNetworkError):
+            client.do_api_query({}, "v1/messages", method="GET")
+
+        # The session should NOT have been closed, since we never
+        # connected in the first place.
+        mock_session.close.assert_not_called()
diff --git a/zulip/zulip/__init__.py b/zulip/zulip/__init__.py
index 6d725a10b..5f7b9428e 100644
--- a/zulip/zulip/__init__.py
+++ b/zulip/zulip/__init__.py
@@ -600,9 +600,6 @@ def do_api_query(
 
         req_files = [(f.name, f) for f in files]
 
-        self.ensure_session()
-        assert self.session is not None
-
         query_state: Dict[str, Any] = {
             "had_error_retry": False,
             "request": request,
@@ -637,6 +634,9 @@ def end_error_retry(succeeded: bool) -> None:
                     print("Failed!")
 
         while True:
+            self.ensure_session()
+            assert self.session is not None
+
             try:
                 kwarg = "params" if method == "GET" else "data"
 
@@ -687,6 +687,14 @@ def end_error_retry(succeeded: bool) -> None:
                         "cannot connect to server " + self.base_url
                     ) from e
 
+                # The session may be holding a stale TCP connection
+                # that was silently dropped by a network middlebox.
+                # Close it so ensure_session() will open a fresh
+                # connection on the next attempt.
+                if self.session is not None:
+                    self.session.close()
+                self.session = None
+
                 if error_retry(""):
                     continue
                 end_error_retry(False)