Fix the reset_client not working issue.

Author: beacoder

python_agent_harness/agent.py

@@ -635,6 +635,13 @@ def safe_delta(text: str) -> None:
                 # sub-agents must not stream into the parent's live
                 # stream row — their text is private until returned
                 on_delta=(safe_delta if self.top_level else None),
+                # a connection error mid-stream discards the partial
+                # output and retries on a fresh client: tell the TUI to
+                # clear the partial text and show that the request is
+                # being restarted
+                on_retry=(
+                    (lambda: session.notify("retry")) if self.top_level else None
+                ),
                 # poll cancellation during retry backoff so Ctrl-C
                 # aborts promptly instead of after the full sleep
                 cancel_check=self._is_cancelled,

python_agent_harness/client.py

@@ -171,6 +171,11 @@ def __init__(
             config.API_RETRY_MAX_DELAY if retry_max_delay is None else retry_max_delay
         )
         self._http = httpx.Client(timeout=timeout, verify=self.verify)
+        # True while the in-flight request was aborted (Ctrl-C): a
+        # connection error on an aborted request must NOT be retried —
+        # the user asked to stop.  Cleared at the start of each chat()
+        # so a fresh turn may retry normally.
+        self._aborted = False
         self.log_path: Path | None = _llm_log_path() if config.LLM_LOG_ENABLED else None
 
     def close(self) -> None:
@@ -188,6 +193,7 @@ def abort(self) -> None:
         the loop treats as a cancel) and then close the pool.  A fresh
         client is swapped in for the next request.
         """
+        self._aborted = True
         old = self._http
         self._http = httpx.Client(timeout=self.timeout, verify=self.verify)
         try:
@@ -202,9 +208,9 @@ def abort(self) -> None:
     def _reset_http(self) -> None:
         """Replace the httpx client with a fresh instance.
 
-        Called after connection-level retries are exhausted so a
-        poisoned pool (stale/dead connections) does not doom every
-        subsequent request in the session.
+        Called before every connection-error retry (and on exhaustion)
+        so a poisoned pool (stale/dead connections) never dooms the
+        retry itself or every subsequent request in the session.
         """
         old = self._http
         self._http = httpx.Client(timeout=self.timeout, verify=self.verify)
@@ -267,6 +273,7 @@ def chat(
         on_tool_call: Callable[[str, str, str], None] | None = None,
         stream: bool = True,
         cancel_check: Callable[[], bool] | None = None,
+        on_retry: Callable[[], None] | None = None,
     ) -> tuple[Message, Usage]:
         """Send a chat request, return (assistant msg, usage).
 
@@ -279,18 +286,27 @@ def chat(
 
         Transient failures (HTTP 429 / 5xx, connection errors) are
         retried with exponential backoff + jitter up to ``retry_max``
-        attempts, honoring ``Retry-After`` when present.  A retry only
-        happens before any delta has been delivered to the callbacks,
-        so streaming output is never duplicated for the caller.  Other
-        4xx errors are permanent and fail immediately.  ``cancel_check``
-        (when given) is polled during backoff sleeps so an abort lands
-        promptly instead of after the full wait.
+        attempts, honoring ``Retry-After`` when present.  A connection
+        error always swaps in a fresh httpx client first (``_reset_http``)
+        so a dead connection never poisons the retry — and a stream that
+        died mid-body IS retried even when deltas already reached the
+        callers: the partial stream is discarded on retry (``on_retry``
+        lets the caller drop its live text), so nothing is duplicated in
+        the returned message.  Other 4xx errors are permanent and fail
+        immediately.  ``cancel_check`` (when given) is polled during
+        backoff sleeps so an abort lands promptly instead of after the
+        full wait.  ``on_retry`` (when given) is invoked right before a
+        retry after a connection error, so a UI can clear the partial
+        output and show that the request is being restarted.
         """
         payload = self._payload(
             messages, tools, stream=stream, temperature=temperature,
             max_tokens=max_tokens, system=system,
             reasoning_effort=reasoning_effort,
         )
+        # a fresh turn may retry connection errors even if a previous
+        # in-flight request was aborted (see abort/_aborted)
+        self._aborted = False
         usage = Usage()
         emitted = False
 
@@ -326,19 +342,22 @@ def wrap_tool_call(name: str, call_id: str, fragment: str) -> None:
                     raise
             except httpx.HTTPError as e:
                 # connection-level failures: connect errors, timeouts,
-                # dropped streams — all transient unless a delta already
-                # reached the caller (then a retry would duplicate it)
-                if emitted or attempt >= self.retry_max:
-                    # Replace the client so a poisoned connection pool
-                    # does not doom all subsequent requests in this
-                    # session.  Without this, a single network hiccup
-                    # can leave the session stuck in a permanent error
-                    # state (the dead connection stays in the pool and
-                    # keeps getting reused).
-                    self._reset_http()
+                # dropped streams.  Swap in a fresh client immediately —
+                # a dead connection must not stay in the pool for the
+                # retry — then retry the request, even when deltas
+                # already reached the caller: the partial stream is
+                # discarded on retry (on_retry lets the caller clear
+                # its live text), so nothing is duplicated in the
+                # stored message.  Only give up once the per-request
+                # attempt budget is exhausted — or immediately when
+                # the request was aborted (Ctrl-C: the user asked to
+                # stop, so a fresh attempt must not be started).
+                self._reset_http()
+                if self._aborted or attempt >= self.retry_max:
                     raise ApiError(f"network error: {e}") from e
+                if on_retry is not None:
+                    on_retry()
                 if self._sleep_backoff(attempt, None, cancel_check):
-                    self._reset_http()
                     raise ApiError(f"network error: {e}") from e
 
         content = "".join(content_parts)

python_agent_harness/tui.py

@@ -423,6 +423,14 @@ def _on_notify(self, kind: str, data: Any = None) -> None:
         elif kind == "compact":
             self.status = " compacted"
             self._history_dirty = True
+        elif kind == "retry":
+            # A connection error mid-stream: the client discarded the
+            # partial response and is retrying on a fresh connection —
+            # drop the partial stream text so the restarted stream
+            # doesn't duplicate it on screen.
+            with self.lock:
+                self.stream_text = ""
+            self.status = " connection lost — retrying"
         elif kind == "todos":
             # TodoWrite updated the task list: the cached history rows
             # (which include the Todos panel) must be rebuilt

tests/test_agent.py

@@ -29,7 +29,7 @@ def __init__(self, script):
 
     def chat(self, messages, tools=None, system=None, temperature=None,
              max_tokens=None, reasoning_effort=None, on_delta=None, stream=True,
-             cancel_check=None):
+             cancel_check=None, on_retry=None):
         self.calls.append([m.to_api() for m in messages])
         self.kwargs.append({
             "tools": tools, "system": system, "temperature": temperature,

tests/test_client.py

@@ -814,8 +814,10 @@ def test_connect_error_non_stream_retried_then_raises(self):
 
         c = make_client(retry_max=2, retry_base_delay=0.01, retry_max_delay=0.05)
         try:
+            # patch the class so the fresh client created by
+            # _reset_http between attempts is covered too
             with mock.patch.object(
-                c._http, "post", side_effect=httpx.ConnectError("refused")
+                httpx.Client, "post", side_effect=httpx.ConnectError("refused")
             ):
                 with self.assertRaises(ApiError):
                     c.chat([Message(role="user", content="hi")], stream=False)
@@ -835,46 +837,98 @@ def test_budget_exhausted_immediately_raises(self):
         finally:
             c.close()
 
-    def test_error_after_delta_not_retried(self):
-        """Once a delta reached the caller, a retry would duplicate it:
-        a stream dying mid-body raises ApiError immediately (no backoff
-        sleep, no second request)."""
+    def test_error_after_delta_retried_on_fresh_client(self):
+        """A stream dying mid-body after deltas is retried on a fresh
+        client: the partial stream is discarded, the retried response
+        is returned, and nothing is duplicated in the stored message."""
+        c = make_client(retry_max=2, retry_base_delay=0.01, retry_max_delay=0.01)
+        old = c._http
+        try:
+            state = {"n": 0}
+
+            def flaky(self_, *a, **kw):
+                state["n"] += 1
+                if state["n"] == 1:
+                    return DieAfterDelta()
+                resp = FakeStreamResp([
+                    'data: {"choices": [{"delta": {"content": "full"}}]}',
+                    "data: [DONE]",
+                ])
+                return FakeStreamCM(resp)
+
+            # patch the class so the fresh client created by
+            # _reset_http is covered too
+            with mock.patch.object(httpx.Client, "stream", flaky):
+                retries = []
+                deltas = []
+                msg, _ = c.chat(
+                    [Message(role="user", content="hi")],
+                    on_delta=deltas.append,
+                    on_retry=lambda: retries.append(1),
+                )
+            self.assertEqual(msg.content, "full")
+            self.assertEqual(state["n"], 2)
+            self.assertEqual(retries, [1])
+            # the retry ran on a fresh (reset) client
+            self.assertIsNot(c._http, old)
+            # the caller saw the partial delta, but the stored message
+            # only carries the retried response
+            self.assertEqual(deltas, ["partial", "full"])
+        finally:
+            c.close()
+
+    def test_error_after_delta_retries_exhausted_then_raises(self):
+        """A stream dying mid-body after deltas retries (reset client
+        each time); when the attempt budget is exhausted it raises
+        ApiError — with the client left fresh for the next call."""
         from python_agent_harness.client import ApiError
 
-        c = make_client(retry_max=3, retry_base_delay=60.0, retry_max_delay=60.0)
+        c = make_client(retry_max=2, retry_base_delay=0.01, retry_max_delay=0.01)
+        old = c._http
         try:
-            class DieAfterDelta:
-                def __init__(self):
-                    self._sent = False
+            with mock.patch.object(
+                httpx.Client, "stream", return_value=DieAfterDelta()
+            ):
+                retries = []
+                with self.assertRaises(ApiError):
+                    c.chat(
+                        [Message(role="user", content="hi")],
+                        on_retry=lambda: retries.append(1),
+                    )
+            self.assertEqual(retries, [1])  # one retry after the first drop
+            self.assertIsNot(c._http, old)  # pool was reset
+            self.assertFalse(c._http.is_closed)
+        finally:
+            c.close()
 
-                def __enter__(self):
-                    return self
 
-                def __exit__(self, *exc):
-                    return False
+class DieAfterDelta:
+    """Streaming response that yields one delta then dies with a
+    connection error (mimics a connection reset mid-body)."""
 
-                @property
-                def status_code(self):
-                    return 200
+    def __init__(self):
+        self._sent = False
 
-                @property
-                def headers(self):
-                    return {}
+    def __enter__(self):
+        return self
 
-                def read(self):
-                    return b""
+    def __exit__(self, *exc):
+        return False
 
-                def iter_lines(self):
-                    yield 'data: {"choices": [{"delta": {"content": "partial"}}]}'
-                    raise httpx.ReadError("connection reset")
+    @property
+    def status_code(self):
+        return 200
 
-            with mock.patch.object(c._http, "stream", return_value=DieAfterDelta()):
-                deltas = []
-                with self.assertRaises(ApiError):
-                    c.chat([Message(role="user", content="hi")], on_delta=deltas.append)
-            self.assertEqual(deltas, ["partial"])
-        finally:
-            c.close()
+    @property
+    def headers(self):
+        return {}
+
+    def read(self):
+        return b""
+
+    def iter_lines(self):
+        yield 'data: {"choices": [{"delta": {"content": "partial"}}]}'
+        raise httpx.ReadError("connection reset")
 
 
 class TestClientResetHttp(unittest.TestCase):

tests/test_subagent.py

@@ -20,7 +20,7 @@ def __init__(self):
 
     def chat(self, messages, tools=None, system=None, temperature=None,
               max_tokens=None, reasoning_effort=None, on_delta=None, stream=True,
-              cancel_check=None):
+              cancel_check=None, on_retry=None):
         self.systems.append(system)
         return Message(role="assistant", content="sub-agent done"), Usage(input_tokens=10)
 

tests/test_subagent_isolation.py

@@ -24,7 +24,7 @@ def __init__(self, script):
 
     def chat(self, messages, tools=None, system=None, temperature=None,
              max_tokens=None, reasoning_effort=None, on_delta=None, stream=True,
-             cancel_check=None):
+             cancel_check=None, on_retry=None):
         self.n += 1
         self.sent.append([m.to_api() for m in messages])
         self.sent_tools.append([t.name for t in tools] if tools else None)

tests/test_todos_scope.py

@@ -22,7 +22,7 @@ def __init__(self, sub_todos):
 
     def chat(self, messages, tools=None, system=None, temperature=None,
              max_tokens=None, reasoning_effort=None, on_delta=None, stream=True,
-             cancel_check=None):
+             cancel_check=None, on_retry=None):
         self.n += 1
         self.sent_tools.append([t.name for t in tools] if tools else None)
         if self.n == 1: