🪲 BUG-#79 #80: Fix message loss in long interactions (provider, chat, serve, agent)

docs/examples/provider.example.json

@@ -5,7 +5,7 @@
   "headers": {
     "X-Org": "acme"
   },
-  "timeout": 60,
+  "timeout": 180,
 
   "_comment": "response_paths is optional. Omit it entirely if the API already returns the OpenAI chat-completions shape (choices[0].message.content, usage.prompt_tokens, ...). Keep it only to remap a different shape, like the example below.",
   "response_paths": {

pycodeloop/cli/chat.py

@@ -431,6 +431,20 @@ async def _run_turn(
             )
         except Exception as exc:
             self.call_from_thread(self._stop_thinking)
+            if self._text_buffer.strip():
+                self.call_from_thread(
+                    self._log,
+                    Panel(
+                        Markdown(self._text_buffer),
+                        border_style="grey50",
+                        title="[dim]interrupted[/dim]",
+                        subtitle=(
+                            "[bold white on grey30] Agent [/bold white on grey30]"
+                        ),
+                        subtitle_align="right",
+                    ),
+                )
+                self._text_buffer = ""
             self.call_from_thread(
                 self._log,
                 self._styled("[bold white]✗ Error:[/bold white] ", str(exc)),

pycodeloop/cli/serve.py

@@ -31,6 +31,9 @@
     response,
 )
 
+_HEARTBEAT_INTERVAL = 15.0
+_CONFIRM_TIMEOUT = 120.0
+
 
 class RpcServer:
     """Wires `Agent` callbacks to JSON-RPC notifications instead of the
@@ -52,13 +55,28 @@ def __init__(
         self._confirm_waiters: dict[str, queue.Queue] = {}
         self._cancel_event: threading.Event | None = None
         self._chat_thread: threading.Thread | None = None
+        self._disconnected = False
         self._wire_callbacks()
 
     def _send(self, message: dict) -> None:
+        """Best-effort write of one NDJSON line to stdout. The client
+        (editor extension) can disconnect mid-turn — closing its end of
+        the pipe — at any point, including while a background thread is
+        still streaming deltas for an in-flight turn. Once that happens
+        every further write raises the same broken-pipe error, so this
+        marks the server disconnected and gives up quietly instead of
+        raising out of a callback (which would otherwise abort whatever
+        turn/tool loop is in progress) or crashing a second time from
+        inside an error handler that itself calls `_send`."""
+        if self._disconnected:
+            return
         line = json.dumps(message)
-        with self._out_lock:
-            sys.stdout.write(line + "\n")
-            sys.stdout.flush()
+        try:
+            with self._out_lock:
+                sys.stdout.write(line + "\n")
+                sys.stdout.flush()
+        except OSError:
+            self._disconnected = True
 
     def _notify(self, method: str, params: dict) -> None:
         self._send(notification(method, params))
@@ -137,7 +155,10 @@ def confirm(name: str, preview: str) -> bool | str:
                 {"id": request_id, "name": name, "preview": preview},
             )
             try:
-                return answer_queue.get()
+                return answer_queue.get(timeout=_CONFIRM_TIMEOUT)
+            except queue.Empty:
+                self._notify("chat/confirmTimeout", {"id": request_id})
+                return False
             finally:
                 self._confirm_waiters.pop(request_id, None)
 
@@ -153,8 +174,23 @@ def confirm(name: str, preview: str) -> bool | str:
         agent.on_compact_end = on_compact_end
         agent.confirm = confirm
 
+    def _run_heartbeat(self, stop: threading.Event) -> None:
+        """Emits `chat/heartbeat` every `_HEARTBEAT_INTERVAL` seconds
+        while a turn is in flight. Long reasoning or a long-running
+        tool can otherwise leave the client with no message at all for
+        minutes; a client with its own read timeout may then conclude
+        the process died and drop the connection, losing the turn even
+        though the server was still working on it."""
+        while not stop.wait(_HEARTBEAT_INTERVAL):
+            self._notify("chat/heartbeat", {})
+
     def _run_chat(self, request_id, params: dict) -> None:
         self._cancel_event = threading.Event()
+        heartbeat_stop = threading.Event()
+        heartbeat_thread = threading.Thread(
+            target=self._run_heartbeat, args=(heartbeat_stop,), daemon=True
+        )
+        heartbeat_thread.start()
         try:
             result = self.flow.run(
                 params.get("prompt", ""),
@@ -165,6 +201,8 @@ def _run_chat(self, request_id, params: dict) -> None:
             self._respond(request_id, {"text": result})
         except Exception as exc:
             self._respond_error(request_id, SERVER_ERROR, str(exc))
+        finally:
+            heartbeat_stop.set()
 
     def _run_ask(self, request_id, params: dict) -> None:
         try:
@@ -259,7 +297,11 @@ def serve_forever(self) -> None:
                 continue
             try:
                 request = json.loads(line)
-            except json.JSONDecodeError:
+            except json.JSONDecodeError as exc:
+                console.print(
+                    f"[dim]⚠ dropped malformed request line ({exc}): "
+                    f"{line[:200]!r}[/dim]"
+                )
                 continue
             self.handle(request)
 

pycodeloop/core/agent.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import contextlib
 import threading
 import time
 from collections.abc import Callable
@@ -116,7 +117,27 @@ def __init__(
 
     def _trace(self, event_type: str, **fields) -> None:
         if self.on_trace_event:
-            self.on_trace_event({"type": event_type, **fields})
+            with contextlib.suppress(Exception):
+                self.on_trace_event({"type": event_type, **fields})
+
+    def _safe_call(self, callback: Callable | None, *args) -> None:
+        """Invokes a consumer-supplied `on_*` callback (UI rendering,
+        storage persistence, etc.) without letting a bug on that side
+        abort the turn/tool loop still in progress. Before this, an
+        exception from e.g. `on_message` (a storage write failing) or
+        `on_tool_result` (a rendering bug) propagated straight out of
+        `Agent.run()`/`_run_tool_calls()`, killing the rest of the turn
+        over what should have been a self-contained side effect."""
+        if callback is None:
+            return
+        try:
+            callback(*args)
+        except Exception as exc:
+            self._trace(
+                "callback_error",
+                callback=getattr(callback, "__name__", repr(callback)),
+                error=str(exc),
+            )
 
     def _complete(self, **kwargs) -> ProviderResponse:
         """`provider.complete()` with retry + exponential backoff on
@@ -155,8 +176,7 @@ def _complete(self, **kwargs) -> ProviderResponse:
                         delay=delay,
                         error=str(exc),
                     )
-                    if self.on_retry:
-                        self.on_retry(attempt + 1, delay, exc)
+                    self._safe_call(self.on_retry, attempt + 1, delay, exc)
                     time.sleep(delay)
                     delay *= 2
 
@@ -172,8 +192,7 @@ def _notify_message(self) -> None:
         caller persist incrementally instead of only after the whole
         (possibly long, multi-tool-call) turn finishes, so a crash
         mid-turn doesn't lose everything already done in it."""
-        if self.on_message:
-            self.on_message()
+        self._safe_call(self.on_message)
 
     def _tool_schemas(self) -> list[dict]:
         return [tool.schema() for tool in self.tools.values()]
@@ -252,8 +271,7 @@ def _run_tool_calls(
     ) -> None:
         for call in calls:
             self._trace("tool_call", name=call.name, arguments=call.arguments)
-            if self.on_tool_call:
-                self.on_tool_call(call.name, call.arguments)
+            self._safe_call(self.on_tool_call, call.name, call.arguments)
 
         if cancel_event and cancel_event.is_set():
             results = {call.id: ("Cancelled by user.", True) for call in calls}
@@ -302,8 +320,9 @@ def _run_tool_calls(
                 is_error=is_error,
                 result_len=len(result_text),
             )
-            if self.on_tool_result:
-                self.on_tool_result(call.name, result_text, is_error)
+            self._safe_call(
+                self.on_tool_result, call.name, result_text, is_error
+            )
             session.add_tool_result(call.id, result_text)
             self._notify_message()
 
@@ -332,8 +351,7 @@ def _compact(self, session: Session) -> None:
         if len(turn_starts) <= _COMPACT_KEEP_RECENT_TURNS:
             return
 
-        if self.on_compact_start:
-            self.on_compact_start()
+        self._safe_call(self.on_compact_start)
 
         before_count = len(history)
         cutoff = turn_starts[-_COMPACT_KEEP_RECENT_TURNS]
@@ -365,8 +383,9 @@ def _compact(self, session: Session) -> None:
         self._trace(
             "compact", before=before_count, after=len(session.messages)
         )
-        if self.on_compact_end:
-            self.on_compact_end(before_count, len(session.messages))
+        self._safe_call(
+            self.on_compact_end, before_count, len(session.messages)
+        )
 
     def run(
         self,
@@ -403,8 +422,9 @@ def run(
                 self._compact(session)
 
             tools = self._tool_schemas()
-            if self.on_request:
-                self.on_request(len(session.history()), len(tools))
+            self._safe_call(
+                self.on_request, len(session.history()), len(tools)
+            )
 
             started_at = time.perf_counter()
             response = self._complete(
@@ -418,13 +438,13 @@ def run(
 
             if response.stop_reason == "cancelled":
                 self.usage = self.usage + response.usage
-                if self.on_usage:
-                    self.on_usage(response.usage, self.usage, elapsed)
+                self._safe_call(
+                    self.on_usage, response.usage, self.usage, elapsed
+                )
                 if response.text.strip():
                     session.add_assistant(response.text)
                     self._notify_message()
-                    if self.on_turn_end:
-                        self.on_turn_end()
+                    self._safe_call(self.on_turn_end)
                 self._trace("run_end", reason="cancelled")
                 return "Cancelled by user."
 
@@ -435,24 +455,25 @@ def run(
                 and empty_retries < _MAX_EMPTY_RESPONSE_RETRIES
             ):
                 self.usage = self.usage + response.usage
-                if self.on_usage:
-                    self.on_usage(response.usage, self.usage, elapsed)
+                self._safe_call(
+                    self.on_usage, response.usage, self.usage, elapsed
+                )
 
                 empty_retries += 1
                 self._trace(
                     "empty_response_retry",
                     model=self.provider.model,
                     attempt=empty_retries,
                 )
-                if self.on_retry:
-                    self.on_retry(
-                        empty_retries,
-                        0.0,
-                        RuntimeError(
-                            f"{self.provider.model} returned an empty "
-                            "response with no tool calls"
-                        ),
-                    )
+                self._safe_call(
+                    self.on_retry,
+                    empty_retries,
+                    0.0,
+                    RuntimeError(
+                        f"{self.provider.model} returned an empty "
+                        "response with no tool calls"
+                    ),
+                )
                 started_at = time.perf_counter()
                 response = self._complete(
                     system_prompt=self.system_prompt,
@@ -465,20 +486,21 @@ def run(
 
                 if response.stop_reason == "cancelled":
                     self.usage = self.usage + response.usage
-                    if self.on_usage:
-                        self.on_usage(response.usage, self.usage, elapsed)
+                    self._safe_call(
+                        self.on_usage, response.usage, self.usage, elapsed
+                    )
                     if response.text.strip():
                         session.add_assistant(response.text)
                         self._notify_message()
-                        if self.on_turn_end:
-                            self.on_turn_end()
+                        self._safe_call(self.on_turn_end)
                     self._trace("run_end", reason="cancelled")
                     return "Cancelled by user."
 
             if not response.text.strip() and not response.tool_calls:
                 self.usage = self.usage + response.usage
-                if self.on_usage:
-                    self.on_usage(response.usage, self.usage, elapsed)
+                self._safe_call(
+                    self.on_usage, response.usage, self.usage, elapsed
+                )
 
                 error_text = (
                     f"{self.provider.model} returned an empty response "
@@ -489,14 +511,12 @@ def run(
                 )
                 session.add_assistant(error_text)
                 self._notify_message()
-                if self.on_turn_end:
-                    self.on_turn_end()
+                self._safe_call(self.on_turn_end)
                 self._trace("run_end", reason="empty_response")
                 return error_text
 
             self.usage = self.usage + response.usage
-            if self.on_usage:
-                self.on_usage(response.usage, self.usage, elapsed)
+            self._safe_call(self.on_usage, response.usage, self.usage, elapsed)
 
             self._trace(
                 "turn",
@@ -509,8 +529,9 @@ def run(
             )
 
             session.update_last_context_tokens(response.usage.input_tokens)
-            if self.on_context:
-                self.on_context(response.usage.input_tokens, context_window)
+            self._safe_call(
+                self.on_context, response.usage.input_tokens, context_window
+            )
 
             tool_calls = [
                 {
@@ -523,8 +544,7 @@ def run(
             ]
             session.add_assistant(response.text, tool_calls=tool_calls or None)
             self._notify_message()
-            if self.on_turn_end:
-                self.on_turn_end()
+            self._safe_call(self.on_turn_end)
 
             if not response.tool_calls:
                 self._trace("run_end", reason="done")