Add timeout and retry logic to Ollama requests

agent/core/llm.py

@@ -21,6 +21,8 @@ class LLMClient:
     def __init__(self, config: LLMConfig):
         self.config = config
         self.base_url = config.base_url.rstrip("/")
+        self.max_retries = 3
+        self.retry_delay = 1.0
 
     def _validate_response(self, response: dict, required_fields: list[str]) -> None:
         """Validate that response contains expected fields.
@@ -76,36 +78,49 @@ def _request(
     ) -> dict:
         url = f"{self.base_url}{endpoint}"
         payload = json.dumps(data).encode("utf-8")
-        req = urllib.request.Request(
-            url,
-            data=payload,
-            headers={"Content-Type": "application/json"},
-            method="POST",
-        )
         effective_timeout = timeout if timeout is not None else self.config.timeout
-        try:
-            resp = urllib.request.urlopen(req, timeout=effective_timeout)
-            if stream:
-                return resp  # Return response object for streaming
-            return json.loads(resp.read().decode("utf-8"))
-        except socket.timeout:
-            raise OllamaError(
-                f"Timeout connecting to Ollama at {self.base_url} "
-                f"(waited {self.config.timeout}s). Make sure Ollama is running."
-            )
-        except urllib.error.URLError as e:
-            if isinstance(e.reason, (socket.timeout, socket.error, ConnectionRefusedError, OSError)):
+
+        for attempt in range(self.max_retries):
+            try:
+                req = urllib.request.Request(
+                    url,
+                    data=payload,
+                    headers={"Content-Type": "application/json"},
+                    method="POST",
+                )
+                resp = urllib.request.urlopen(req, timeout=effective_timeout)
+                if stream:
+                    return resp
+                return json.loads(resp.read().decode("utf-8"))
+            except socket.timeout as e:
+                # A timeout means Ollama is still working, not that the
+                # connection failed - retrying would just queue up more
+                # concurrent work against an already-struggling server.
+                raise OllamaError(
+                    f"Timeout connecting to Ollama at {self.base_url} "
+                    f"(waited {effective_timeout}s). Make sure Ollama is running."
+                ) from e
+            except urllib.error.HTTPError as e:
+                # Must be checked before URLError: HTTPError subclasses it.
+                body = e.read().decode("utf-8", errors="replace")
+                raise OllamaError(f"Ollama API error ({e.code}): {body}") from e
+            except urllib.error.URLError as e:
+                if isinstance(e.reason, (ConnectionRefusedError, socket.gaierror)):
+                    if attempt < self.max_retries - 1:
+                        time.sleep(self.retry_delay * (2 ** attempt))
+                        continue
+                    raise OllamaError(
+                        f"Cannot connect to Ollama at {self.base_url} "
+                        f"after {self.max_retries} attempts. "
+                        f"Make sure Ollama is running: {e}"
+                    ) from e
                 raise OllamaError(
                     f"Cannot connect to Ollama at {self.base_url}. "
                     f"Make sure Ollama is running: {e}"
-                )
-            raise OllamaError(
-                f"Cannot connect to Ollama at {self.base_url}. "
-                f"Make sure Ollama is running: {e}"
-            )
-        except urllib.error.HTTPError as e:
-            body = e.read().decode("utf-8", errors="replace")
-            raise OllamaError(f"Ollama API error ({e.code}): {body}")
+                ) from e
+
+        # Unreachable: every loop iteration either returns or raises above.
+        raise OllamaError(f"Failed to reach Ollama at {self.base_url}.")
 
     def check_health(self) -> bool:
         """Check if Ollama is running and the model is available."""

tests/test_llm.py

@@ -1,8 +1,11 @@
 """Tests for the Ollama LLM client response validation."""
 
+import io
 import json
+import socket
 import sys
 import unittest
+import urllib.error
 from pathlib import Path
 from unittest.mock import patch
 
@@ -166,5 +169,62 @@ def test_non_dict_message_is_skipped(self):
         self.assertEqual(result["message"]["content"], "ok")
 
 
+class TestRequestRetry(unittest.TestCase):
+    """Test the retry behavior in LLMClient._request."""
+
+    def test_connection_refused_is_retried_and_eventually_raises(self):
+        client = make_client()
+        error = urllib.error.URLError(ConnectionRefusedError())
+        with patch("urllib.request.urlopen", side_effect=error) as mock_urlopen, \
+                patch("time.sleep") as mock_sleep:
+            with self.assertRaises(OllamaError):
+                client._request("/api/chat", {})
+        self.assertEqual(mock_urlopen.call_count, client.max_retries)
+        self.assertEqual(mock_sleep.call_count, client.max_retries - 1)
+
+    def test_dns_failure_is_retried_and_eventually_raises(self):
+        client = make_client()
+        error = urllib.error.URLError(socket.gaierror())
+        with patch("urllib.request.urlopen", side_effect=error) as mock_urlopen, \
+                patch("time.sleep"):
+            with self.assertRaises(OllamaError):
+                client._request("/api/chat", {})
+        self.assertEqual(mock_urlopen.call_count, client.max_retries)
+
+    def test_socket_timeout_fails_fast_without_retry(self):
+        """A timeout means Ollama is still generating - retrying would only
+        make a slow response slower, so it must not be retried."""
+        client = make_client()
+        with patch("urllib.request.urlopen", side_effect=socket.timeout()) as mock_urlopen, \
+                patch("time.sleep") as mock_sleep:
+            with self.assertRaises(OllamaError):
+                client._request("/api/chat", {})
+        self.assertEqual(mock_urlopen.call_count, 1)
+        mock_sleep.assert_not_called()
+
+    def test_non_connection_url_error_fails_fast_without_retry(self):
+        client = make_client()
+        error = urllib.error.URLError("ssl handshake failed")
+        with patch("urllib.request.urlopen", side_effect=error) as mock_urlopen, \
+                patch("time.sleep") as mock_sleep:
+            with self.assertRaises(OllamaError):
+                client._request("/api/chat", {})
+        self.assertEqual(mock_urlopen.call_count, 1)
+        mock_sleep.assert_not_called()
+
+    def test_http_error_surfaces_response_body(self):
+        """HTTPError subclasses URLError, so its handler must be checked
+        first or the actual API error body never reaches the caller."""
+        client = make_client()
+        http_error = urllib.error.HTTPError(
+            "http://x", 400, "Bad Request", {}, io.BytesIO(b'{"error": "bad model"}')
+        )
+        with patch("urllib.request.urlopen", side_effect=http_error):
+            with self.assertRaises(OllamaError) as ctx:
+                client._request("/api/chat", {})
+        self.assertIn("bad model", str(ctx.exception))
+        self.assertIn("400", str(ctx.exception))
+
+
 if __name__ == "__main__":
     unittest.main()