Merge branch 'main' into fix/tunnel-close-unhandled-1114

Co-authored-by: Isaac

omnigent/server/routes/sessions.py

@@ -3726,7 +3726,7 @@ async def _forward_approval_to_runner(
             json={"type": _APPROVAL_TYPE, "data": data},
             timeout=10.0,
         )
-    except httpx.HTTPError:
+    except (httpx.HTTPError, ConnectionError):
         _logger.exception(
             "Approval forward failed for %r",
             session_id,
@@ -6497,7 +6497,7 @@ async def _proxy_get_session_resources_to_runner(
         )
     except HTTPException:
         raise
-    except httpx.HTTPError as exc:
+    except (httpx.HTTPError, ConnectionError) as exc:
         _logger.warning(
             "session resources: runner call failed for session=%s (%s)",
             session_id,
@@ -7360,7 +7360,7 @@ async def _forward_session_change_to_runner(
             json=event,
             timeout=5.0,
         )
-    except httpx.HTTPError:
+    except (httpx.HTTPError, ConnectionError):
         _logger.exception(
             "Session-change forward failed for session=%r type=%r",
             session_id,
@@ -7686,7 +7686,7 @@ async def _resolve_skill_meta_text_via_runner(
             json={"name": skill_name, "arguments": arguments},
             timeout=10.0,
         )
-    except httpx.HTTPError as exc:
+    except (httpx.HTTPError, ConnectionError) as exc:
         raise OmnigentError(
             f"Runner unreachable while resolving skill {skill_name!r}: {exc}",
             code=ErrorCode.INTERNAL_ERROR,
@@ -7853,7 +7853,7 @@ async def _dispatch_skill_slash_command_to_runner(
         )
         event = OutputItemDoneEvent(type="response.output_item.done", item=visible.to_api_dict())
         session_stream.publish(session_id, event.model_dump())
-    except httpx.HTTPError:
+    except (httpx.HTTPError, ConnectionError):
         _logger.exception(
             "Forward of skill slash command failed for session=%s; "
             "items persisted, runner picks up on reconnect.",
@@ -8139,7 +8139,7 @@ async def _forward_event_to_runner(
         # Publish input.consumed AFTER the forward succeeds —
         # the runner has the message and will start the turn.
         _publish_input_consumed(session_id, persisted_items[0])
-    except httpx.HTTPError:
+    except (httpx.HTTPError, ConnectionError):
         _logger.exception(
             "Forward to runner failed for session=%s; "
             "event persisted, runner picks up on reconnect.",
@@ -9125,10 +9125,20 @@ async def _relay_runner_stream(
         # close; treat the same as HTTPError so the task exits
         # gracefully instead of leaving an unretrieved exception.
         _logger.warning(
-            "Relay: ended for session=%s",
+            "Relay: runner transport lost for session=%s",
             session_id,
             exc_info=True,
         )
+        # Publish a failed status so the client's SSE stream sees a
+        # clean error event instead of silent truncation (#1114).
+        _publish_status(
+            session_id,
+            "failed",
+            ErrorDetail(
+                code="runner_disconnected",
+                message="Runner disconnected unexpectedly.",
+            ),
+        )
     except asyncio.CancelledError:
         raise
     finally:
@@ -11688,7 +11698,7 @@ async def _notify_runner_of_bundled_child(
             },
             timeout=10.0,
         )
-    except httpx.HTTPError:
+    except (httpx.HTTPError, ConnectionError):
         _logger.warning(
             "Failed to notify runner about bundled session %s",
             session_id,
@@ -12801,7 +12811,7 @@ def create_sessions_router(
                     },
                     timeout=10.0,
                 )
-            except httpx.HTTPError:
+            except (httpx.HTTPError, ConnectionError):
                 _logger.warning(
                     "Failed to notify runner about session %s",
                     resp.id,
@@ -15692,7 +15702,7 @@ def create_sessions_router(
             )
         try:
             resp = await runner_client.get(path, params=params, timeout=10.0)
-        except httpx.HTTPError as exc:
+        except (httpx.HTTPError, ConnectionError) as exc:
             raise HTTPException(
                 status_code=502,
                 detail="runner resource endpoint unavailable",
@@ -15739,7 +15749,7 @@ def create_sessions_router(
                 json=body,
                 timeout=10.0,
             )
-        except httpx.HTTPError as exc:
+        except (httpx.HTTPError, ConnectionError) as exc:
             raise HTTPException(
                 status_code=502,
                 detail="runner resource endpoint unavailable",
@@ -15767,7 +15777,7 @@ def create_sessions_router(
             )
         try:
             resp = await runner_client.delete(path, timeout=10.0)
-        except httpx.HTTPError as exc:
+        except (httpx.HTTPError, ConnectionError) as exc:
             raise HTTPException(
                 status_code=502,
                 detail="runner resource endpoint unavailable",
@@ -15801,7 +15811,7 @@ def create_sessions_router(
                 json=body,
                 timeout=10.0,
             )
-        except httpx.HTTPError as exc:
+        except (httpx.HTTPError, ConnectionError) as exc:
             raise HTTPException(
                 status_code=502,
                 detail="runner resource endpoint unavailable",
@@ -15835,7 +15845,7 @@ def create_sessions_router(
                 json=body,
                 timeout=10.0,
             )
-        except httpx.HTTPError as exc:
+        except (httpx.HTTPError, ConnectionError) as exc:
             raise HTTPException(
                 status_code=502,
                 detail="runner resource endpoint unavailable",
@@ -17605,7 +17615,7 @@ def create_sessions_router(
                     },
                     timeout=10.0,
                 )
-            except httpx.HTTPError as exc:
+            except (httpx.HTTPError, ConnectionError) as exc:
                 # Fail loud (503), not best-effort: unlike the advisory
                 # interrupt-forward, a dropped tool_result leaves the parked
                 # turn hanging until it times out. Surfacing the failure lets
@@ -18158,7 +18168,7 @@ def create_sessions_router(
                     f"/v1/sessions/{session_id}/resources",
                     timeout=10.0,
                 )
-            except httpx.HTTPError:
+            except (httpx.HTTPError, ConnectionError):
                 _logger.warning(
                     "Runner cleanup failed for %s, falling back",
                     session_id,
@@ -18920,7 +18930,7 @@ async def _load_runner_skills(
             f"/v1/sessions/{session_id}/skills",
             timeout=5.0,
         )
-    except httpx.HTTPError:
+    except (httpx.HTTPError, ConnectionError):
         _logger.debug("Runner skills query failed for %s", session_id)
         return
     if resp.status_code != 200:
@@ -19006,7 +19016,7 @@ async def _load_model_options(
     for attempt in range(len(_CODEX_MODEL_OPTIONS_RETRY_DELAYS_S) + 1):
         try:
             resp = await runner_client.get(path, timeout=5.0)
-        except httpx.HTTPError:
+        except (httpx.HTTPError, ConnectionError):
             _logger.debug("Runner Codex model-options query failed for %s", session_id)
             return
         if resp.status_code != 200:
@@ -19158,7 +19168,7 @@ async def _get_session_snapshot(
                     raw = resp.json().get("status", "idle")
                     _session_status_cache[session_id] = raw
                     status = _session_status_from_cache(session_id)
-            except httpx.HTTPError:
+            except (httpx.HTTPError, ConnectionError):
                 _logger.debug(
                     "Runner status query failed for %s",
                     session_id,

tests/server/routes/test_sessions_runner_relay.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import json
 from collections.abc import AsyncIterator
 from types import TracebackType
@@ -363,3 +364,108 @@ async def test_relay_text_flush_publishes_persisted_item(db_uri: str) -> None:
             await asyncio.wait_for(handle.task, timeout=1.0)
         sessions_module._runner_relay_tasks.clear()
         session_stream.close(session_id)
+
+
+class _TunnelCloseStreamResponse:
+    """
+    Async context manager that raises ``ConnectionError`` mid-stream.
+
+    Emits the ready heartbeat, waits for a gate, then raises
+    ``ConnectionError`` to simulate a ws-tunnel drop.
+
+    :param gate: Event the test sets once its collector is subscribed,
+        so the error fires after the collector can observe it.
+    """
+
+    def __init__(self, gate: asyncio.Event) -> None:
+        self._gate = gate
+
+    async def __aenter__(self) -> _TunnelCloseStreamResponse:
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        del exc_type, exc, traceback
+
+    async def aiter_text(self) -> AsyncIterator[str]:
+        yield 'data: {"type": "session.heartbeat"}\n\n'
+        await self._gate.wait()
+        raise ConnectionError("tunnel closed before request completed")
+
+
+class _TunnelCloseRunnerClient:
+    """Fake runner client whose stream drops with ``ConnectionError``.
+
+    :param gate: Event that gates the error (set by the test once
+        its stream collector is subscribed).
+    """
+
+    def __init__(self, gate: asyncio.Event) -> None:
+        self._gate = gate
+
+    def stream(
+        self,
+        method: str,
+        path: str,
+        *,
+        timeout: Any,
+    ) -> _TunnelCloseStreamResponse:
+        del method, path, timeout
+        return _TunnelCloseStreamResponse(self._gate)
+
+
+@pytest.mark.asyncio
+async def test_relay_publishes_failed_status_on_tunnel_close() -> None:
+    """
+    A tunnel close mid-stream publishes ``session.status`` "failed".
+
+    Regression test for #1114: before the fix the relay swallowed the
+    ``ConnectionError`` and exited silently, leaving the client's SSE
+    stream truncated with no error event.
+    """
+    from omnigent.runtime import session_stream
+    from omnigent.server.routes import sessions as sessions_module
+
+    sessions_module._runner_relay_tasks.clear()
+    gate = asyncio.Event()
+    fake_runner = _TunnelCloseRunnerClient(gate)
+    session_id = "conv_tunnel_close"
+
+    collector = None
+    try:
+        handle = await sessions_module._ensure_runner_relay_ready(
+            session_id,
+            "runner_tunnel_close",
+            fake_runner,  # type: ignore[arg-type]
+            conversation_store=None,
+        )
+        assert handle is not None
+
+        # Subscribe BEFORE releasing the error so the published
+        # session.status event fans out to the collector.
+        collector = await start_session_stream_collector(session_id)
+        gate.set()
+
+        # The relay task should finish quickly after the ConnectionError.
+        await asyncio.wait_for(handle.task, timeout=2.0)
+
+        # Wait for the failed-status event to arrive at the collector.
+        event = await asyncio.wait_for(collector.queue.get(), timeout=2.0)
+        assert event.get("type") == "session.status"
+        assert event.get("status") == "failed"
+        assert event["error"]["code"] == "runner_disconnected"
+    finally:
+        gate.set()
+        if collector is not None:
+            await collector.stop()
+        handle = sessions_module._runner_relay_tasks.get(session_id)
+        if handle is not None and not handle.task.done():
+            handle.task.cancel()
+            with contextlib.suppress(asyncio.CancelledError, asyncio.TimeoutError):
+                await asyncio.wait_for(handle.task, timeout=1.0)
+        sessions_module._runner_relay_tasks.clear()
+        session_stream.close(session_id)