Keep agent lanes alive after controller response failures

agent/ec_agent.py

@@ -12,6 +12,7 @@ import contextlib
 from datetime import datetime, timezone
 import math
 import hashlib
+import http.client
 import fcntl
 import json
 import os
@@ -32,7 +33,7 @@ import urllib.parse
 import urllib.request
 import uuid
 
-VERSION = "1.5.1"
+VERSION = "1.5.2"
 LABEL = "io.electric-cherry"
 MAX_FILE = 32 * 1024 * 1024
 MANIFEST_MAX_ENTRIES = 10000
@@ -1099,7 +1100,7 @@ class Agent:
                     raise AgentError("Controller redirects are not allowed")
             with urllib.request.build_opener(NoRedirect()).open(req, timeout=35) as response:
                 return json.load(response)
-        except (OSError, ValueError):
+        except (OSError, ValueError, http.client.HTTPException):
             raise AgentError("Controller request failed") from None
 
     def cpu_usage_percent(self):
@@ -2354,6 +2355,8 @@ class Agent:
                 self.heartbeat()
             except (AgentError, OSError, ValueError):
                 pass
+            except Exception as err:
+                print(f"Heartbeat recovered from {type(err).__name__}", file=sys.stderr)
 
     @staticmethod
     def read_operation(operation):
@@ -2421,6 +2424,9 @@ class Agent:
                         self.stopping.wait(0.25)
                 except (AgentError, OSError, ValueError, KeyError, sqlite3.Error):
                     self.stopping.wait(3)
+                except Exception as err:
+                    print(f"{lane} lane recovered from {type(err).__name__}", file=sys.stderr)
+                    self.stopping.wait(3)
         finally:
             journal.close()
 
@@ -2475,6 +2481,9 @@ class Agent:
                 break
             except (AgentError, OSError, ValueError):
                 self.stopping.wait(3)
+            except Exception as err:
+                print(f"Initial heartbeat recovered from {type(err).__name__}", file=sys.stderr)
+                self.stopping.wait(3)
         heartbeats = threading.Thread(target=self.heartbeat_loop, daemon=True)
         heartbeats.start()
         # Preserve two concurrent health checks even while files are being read.
@@ -2488,6 +2497,8 @@ class Agent:
             except (AgentError, OSError, ValueError, KeyError):
                 # Logs deliberately omit URLs, tokens, payloads and raw errors.
                 print("Host controller contact unavailable; existing containers continue running", file=sys.stderr)
+            except Exception as err:
+                print(f"Write lane recovered from {type(err).__name__}", file=sys.stderr)
             self.stopping.wait(0.25 if completed else 3)
         for reader in readers:
             reader.join(timeout=2)

agent/tests/test_agent.py

@@ -1,5 +1,6 @@
 import base64
 import importlib.util
+import http.client
 import io
 import json
 from pathlib import Path
@@ -392,6 +393,59 @@ class AgentTests(unittest.TestCase):
         self.agent.journal.close()
         self.temp.cleanup()
 
+    def test_controller_truncated_response_is_a_sanitized_retryable_error(self):
+        # Exercise Python's actual HTTP body reader against a truncated stream.
+        sock = MagicMock()
+        sock.makefile.return_value = io.BytesIO(b'HTTP/1.1 200 OK\r\nContent-Length: 100\r\n\r\n{"secret":')
+        response = http.client.HTTPResponse(sock)
+        response.begin()
+        with patch.object(module.urllib.request, "build_opener") as opener:
+            opener.return_value.open.return_value = response
+            with self.assertRaisesRegex(AgentError, "^Controller request failed$"):
+                self.agent.request("GET", "/api/agent/operations")
+
+    def test_read_and_health_lanes_retry_unexpected_failures_without_logging_secrets(self):
+        for lane in ("read", "health"):
+            with self.subTest(lane=lane):
+                self.agent.stopping.clear()
+                polls = []
+
+                def request(*args):
+                    polls.append(args)
+                    if len(polls) == 1:
+                        raise RuntimeError("secret controller response")
+                    self.agent.stopping.set()
+                    return {"operations": []}
+
+                with patch.object(self.agent, "request", side_effect=request), patch.object(self.agent.stopping, "wait", return_value=False), patch.object(module.sys, "stderr", new_callable=io.StringIO) as logs:
+                    self.agent.read_loop(lane)
+                self.assertEqual(len(polls), 2)
+                self.assertIn("RuntimeError", logs.getvalue())
+                self.assertNotIn("secret", logs.getvalue())
+
+    def test_heartbeat_loop_survives_unexpected_failures(self):
+        with patch.object(self.agent.stopping, "wait", side_effect=[False, False, True]), patch.object(self.agent, "heartbeat", side_effect=[RuntimeError("secret"), None]) as heartbeat, patch.object(module.sys, "stderr", new_callable=io.StringIO) as logs:
+            self.agent.heartbeat_loop()
+        self.assertEqual(heartbeat.call_count, 2)
+        self.assertIn("RuntimeError", logs.getvalue())
+        self.assertNotIn("secret", logs.getvalue())
+
+    def test_startup_and_write_loop_survive_unexpected_failures(self):
+        polls = []
+
+        def poll():
+            polls.append(True)
+            if len(polls) == 1:
+                raise RuntimeError("secret")
+            self.agent.stopping.set()
+
+        with patch.object(module.threading, "Thread"), patch.object(self.agent.stopping, "wait", return_value=False), patch.object(self.agent, "heartbeat", side_effect=[RuntimeError("secret"), None]) as heartbeat, patch.object(self.agent, "poll_operations", side_effect=poll), patch.object(module.sys, "stderr", new_callable=io.StringIO) as logs:
+            self.agent.run()
+        self.assertEqual(heartbeat.call_count, 2)
+        self.assertEqual(len(polls), 2)
+        self.assertEqual(logs.getvalue().count("RuntimeError"), 2)
+        self.assertNotIn("secret", logs.getvalue())
+
     def test_inventory_reuses_manifests_until_mtime_changes(self):
         marker = self.library / "manifest.json"
         saved = json.loads(marker.read_text())