fix: preserve edge streaming responses on upstream failure

README.md

@@ -157,6 +157,8 @@ When `services.olla.version` is omitted inside an explicit `services.olla` block
 
 After changing model placement or node topology in `tfconfig.yaml`, run `make restart gateway-cache-01` (or full `make restart`) before `make smoke <node>` so Olla and TF edge reload the generated router config.
 
+TF edge defaults to a 180-second upstream I/O timeout, leaving room for the generated Olla config's 120-second backend timeouts and connection attempts. This limits waits for network activity, not total generation time. If an upstream stream fails after HTTP headers have been sent, edge closes the incomplete response and records a 502 failure in its access log; it cannot change the status already sent to the client. Streaming bodies are captured only when `services.edge.inspector.enabled` is true. Inspector response capture is limited to the first 1 MiB, with `body_truncated: true` in the JSONL entry when clipped; clients still receive the full stream.
+
 ### GPU memory ceiling on inference nodes
 
 Set `services.omlx.gpu_memory_gb` in `tfconfig.yaml` for a shared ceiling, or override it with `nodes.<node>.runtime.gpu_memory_gb`. Use a positive integer in **GiB**. Omitted node values inherit the service setting; explicit `null` selects actual RAM minus 10 GiB. A 128 GiB node therefore defaults to 118 GiB (`iogpu.wired_limit_mb=120832`). RAM comes from `hw.memsize` on the inference node, not the controller or the inventory's `ram_gb`.

src/thunder_forge/cluster/edge.py

@@ -61,6 +61,7 @@
 EDGE_DEFAULT_PORT = DEFAULT_EDGE_PORT
 EDGE_LAUNCHD_LABEL_PREFIX = "com.thunder-forge.edge"
 EDGE_DEFAULT_MAX_BODY_BYTES = 16 * 1024 * 1024
+EDGE_MAX_STREAM_CAPTURE_BYTES = 1024 * 1024
 
 _ENV_LINE_RE = re.compile(
     r"^(?P<prefix>\s*)(?P<name>[A-Za-z_][A-Za-z0-9_]*)(?P<sep>\s*=\s*)(?P<value>.*?)(?P<suffix>\s*)$"
@@ -106,7 +107,8 @@ class EdgeProxyConfig:
     active_requests: dict[str, tuple[str, str]] = field(default_factory=dict)
     active_requests_lock: threading.Lock = field(default_factory=threading.Lock, repr=False)
     model_catalog: list[EdgeModelCatalogEntry] = field(default_factory=list)
-    timeout: float = 60.0
+    # Leave room for Olla's 120-second backend timeouts and connection attempts.
+    timeout: float = 180.0
     status_timeout: float = 8.0
     max_body_bytes: int = EDGE_DEFAULT_MAX_BODY_BYTES
     inspector: EdgeInspectorConfig | None = None
@@ -1166,6 +1168,7 @@ def _write_edge_inspector_entry(
     status_code: int | None = None,
     headers: dict[str, str] | httpx.Headers | None = None,
     body: bytes | None = None,
+    body_truncated: bool = False,
     error: str | None = None,
 ) -> None:
     inspector = config.inspector
@@ -1188,6 +1191,8 @@ def _write_edge_inspector_entry(
         payload["headers"] = _normalize_headers_for_dump(headers)
     if body is not None:
         payload["body"] = _body_dump_payload(body)
+    if body_truncated:
+        payload["body_truncated"] = True
     if error:
         payload["error"] = error
 
@@ -1574,9 +1579,10 @@ def record_stream(status_code: int, *, olla_endpoint: str = "") -> None:
                     olla_endpoint=olla_endpoint,
                 )
                 logged = True
-                _unregister_active_request(config, upstream.request_id)
 
             normalized_base_url = config.olla_base_url.rstrip("/")
+            headers_sent = False
+            olla_endpoint = ""
             try:
                 with httpx.Client(base_url=normalized_base_url, timeout=config.timeout, trust_env=False) as client:
                     with client.stream(
@@ -1585,7 +1591,8 @@ def record_stream(status_code: int, *, olla_endpoint: str = "") -> None:
                         headers=upstream.forwarded_headers,
                         content=body,
                     ) as response:
-                        captured_body = bytearray()
+                        captured_body = bytearray() if inspector is not None and inspector.enabled else None
+                        body_truncated = False
                         olla_endpoint = response.headers.get("X-Olla-Endpoint", "")
                         self.send_response(response.status_code)
                         response_headers = _proxy_response_headers(response.headers)
@@ -1594,21 +1601,25 @@ def record_stream(status_code: int, *, olla_endpoint: str = "") -> None:
                             self.send_header(name, value)
                         try:
                             self.end_headers()
+                            headers_sent = True
                         except OSError as exc:
                             if self._is_client_disconnect(exc):
                                 record_stream(response.status_code, olla_endpoint=olla_endpoint)
                                 return
                             raise
                         for chunk in response.iter_bytes():
                             if chunk:
-                                captured_body.extend(chunk)
+                                if captured_body is not None:
+                                    remaining = EDGE_MAX_STREAM_CAPTURE_BYTES - len(captured_body)
+                                    captured_body.extend(chunk[:remaining])
+                                    body_truncated |= len(chunk) > remaining
                                 try:
                                     self.wfile.write(chunk)
                                     self.wfile.flush()
                                 except OSError as exc:
                                     if self._is_client_disconnect(exc):
                                         record_stream(response.status_code, olla_endpoint=olla_endpoint)
-                                        if inspector is not None and inspector.enabled:
+                                        if captured_body is not None:
                                             _write_edge_inspector_entry(
                                                 config,
                                                 event_type="response",
@@ -1618,10 +1629,11 @@ def record_stream(status_code: int, *, olla_endpoint: str = "") -> None:
                                                 status_code=response.status_code,
                                                 headers=response.headers,
                                                 body=bytes(captured_body),
+                                                body_truncated=body_truncated,
                                             )
                                         return
                                     raise
-                        if inspector is not None and inspector.enabled:
+                        if captured_body is not None:
                             _write_edge_inspector_entry(
                                 config,
                                 event_type="response",
@@ -1631,6 +1643,7 @@ def record_stream(status_code: int, *, olla_endpoint: str = "") -> None:
                                 status_code=response.status_code,
                                 headers=response.headers,
                                 body=bytes(captured_body),
+                                body_truncated=body_truncated,
                             )
                         record_stream(response.status_code, olla_endpoint=olla_endpoint)
             except httpx.HTTPError as exc:
@@ -1643,8 +1656,13 @@ def record_stream(status_code: int, *, olla_endpoint: str = "") -> None:
                         path=self.path,
                         error=str(exc),
                     )
-                record_stream(502)
-                self._send_edge_response(_json_response(502, {"error": f"upstream_failed: {exc}"}))
+                record_stream(502, olla_endpoint=olla_endpoint)
+                if headers_sent:
+                    self.close_connection = True
+                else:
+                    self._send_edge_response(_json_response(502, {"error": f"upstream_failed: {exc}"}))
+            finally:
+                _unregister_active_request(config, upstream.request_id)
 
         def _send_edge_response(self, result: EdgeProxyResponse) -> None:
             self.send_response(result.status_code)

tests/test_edge_streaming.py

@@ -0,0 +1,182 @@
+"""Exercise the actual streaming HTTP handler over a local socket pair."""
+
+import json
+import socket
+import threading
+import time
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from types import SimpleNamespace
+
+import httpx
+import pytest
+
+from thunder_forge.cluster import edge
+
+
+@pytest.fixture
+def exchange(monkeypatch):
+    client_type = httpx.Client
+
+    def run(config, upstream, *, timeout_scale=1):
+        handlers = []
+
+        def server(_address, handler):
+            handlers.append(handler)
+            return SimpleNamespace(serve_forever=lambda: None)
+
+        monkeypatch.setattr(edge, "ThreadingHTTPServer", server)
+        def proxy_client(**kwargs):
+            kwargs["timeout"] *= timeout_scale
+            return client_type(**kwargs, transport=httpx.MockTransport(upstream) if upstream else None)
+
+        monkeypatch.setattr(edge.httpx, "Client", proxy_client)
+        edge.serve_edge_proxy(host="127.0.0.1", port=0, config=config)
+        body = json.dumps({"model": "agent", "messages": [], "stream": True}).encode()
+        request = (
+            b"POST /v1/chat/completions HTTP/1.0\r\n"
+            b"Authorization: Bearer fixture-key\r\n"
+            + f"Content-Length: {len(body)}\r\n\r\n".encode() + body
+        )
+        errors = []
+        client, connection = socket.socketpair()
+
+        def handle():
+            try:
+                handlers[0](connection, ("127.0.0.1", 1234), None)
+            except Exception as exc:
+                errors.append(exc)
+            finally:
+                connection.close()
+
+        with client:
+            client.settimeout(5)
+            client.sendall(request)
+            client.shutdown(socket.SHUT_WR)
+            thread = threading.Thread(target=handle, daemon=True)
+            thread.start()
+            try:
+                with client.makefile("rb") as response:
+                    wire = response.read()
+            finally:
+                thread.join(timeout=5)
+        assert not thread.is_alive()
+        assert not errors
+        return wire
+
+    return run
+
+
+def test_edge_allows_a_response_within_olla_timeout(exchange):
+    # Scale 90 seconds of header latency to 450 ms: Olla allows 120 seconds.
+    class Upstream(BaseHTTPRequestHandler):
+        def do_POST(self):
+            self.rfile.read(int(self.headers["Content-Length"]))
+            time.sleep(0.45)
+            try:
+                self.send_response(200)
+                self.send_header("Content-Type", "text/event-stream")
+                self.end_headers()
+                self.wfile.write(b"data: [DONE]\n\n")
+            except (BrokenPipeError, ConnectionResetError):
+                pass
+
+        def log_message(self, *_args):
+            pass
+
+    with ThreadingHTTPServer(("127.0.0.1", 0), Upstream) as upstream:
+        thread = threading.Thread(target=upstream.serve_forever, kwargs={"poll_interval": 0.01}, daemon=True)
+        thread.start()
+        config = edge.EdgeProxyConfig(
+            olla_base_url=f"http://127.0.0.1:{upstream.server_port}",
+            clients_by_key={"fixture-key": edge.EdgeClient("fixture")},
+        )
+        try:
+            wire = exchange(config, None, timeout_scale=0.005)
+        finally:
+            upstream.shutdown()
+            thread.join(timeout=5)
+    assert wire.startswith(b"HTTP/1.0 200")
+    assert wire.endswith(b"data: [DONE]\n\n")
+
+
+@pytest.mark.parametrize("failure", ["connect", "before_body", "after_chunk", None])
+def test_stream_failure_sends_only_one_http_response(exchange, failure):
+    chunk = b'data: {"choices":[{"delta":{"content":"hello"}}]}\n\n'
+    closed = []
+    records = []
+
+    class Stream(httpx.SyncByteStream):
+        def __iter__(self):
+            if failure == "before_body":
+                raise httpx.ReadError("fixture read failure")
+            yield chunk
+            if failure == "after_chunk":
+                raise httpx.ReadError("fixture read failure")
+            yield b"data: [DONE]\n\n"
+
+        def close(self):
+            closed.append(True)
+
+    def upstream(request):
+        if failure == "connect":
+            raise httpx.ConnectError("fixture connection failure", request=request)
+        return httpx.Response(200, headers={
+            "Content-Type": "text/event-stream", "X-Olla-Endpoint": "infer-01",
+        }, stream=Stream())
+
+    config = edge.EdgeProxyConfig(
+        olla_base_url="http://olla.test", clients_by_key={"fixture-key": edge.EdgeClient("fixture")},
+        access_log_sink=records.append,
+    )
+    wire = exchange(config, upstream)
+    assert wire.count(b"HTTP/1.0 ") == 1
+    assert wire.startswith(b"HTTP/1.0 502" if failure == "connect" else b"HTTP/1.0 200")
+    body = wire.split(b"\r\n\r\n", 1)[1]
+    if failure == "connect":
+        assert json.loads(body)["error"].startswith("upstream_failed:")
+    else:
+        assert closed == [True]
+        assert body == (b"" if failure == "before_body" else chunk if failure else chunk + b"data: [DONE]\n\n")
+    assert config.active_requests == {}
+    assert len(records) == 1
+    record = json.loads(records[0])
+    assert record["status_code"] == (502 if failure else 200)
+    if failure != "connect":
+        assert record["olla_endpoint"] == "infer-01"
+
+
+@pytest.mark.parametrize("inspect", [False, True])
+def test_stream_capture_is_optional_and_bounded(exchange, monkeypatch, tmp_path, inspect):
+    chunk = b"data: " + b"x" * 65536 + b"\n\n"
+    chunks = 20
+    retained = []
+
+    class Capture(bytearray):
+        def extend(self, data):
+            super().extend(data)
+            retained.append(len(self))
+
+    class Stream(httpx.SyncByteStream):
+        def __iter__(self):
+            for _ in range(chunks):
+                yield chunk
+
+    monkeypatch.setattr(edge, "bytearray", Capture, raising=False)
+    config = edge.EdgeProxyConfig(
+        olla_base_url="http://olla.test", clients_by_key={"fixture-key": edge.EdgeClient("fixture")},
+        inspector=edge.EdgeInspectorConfig(enabled=inspect, output_dir=tmp_path),
+    )
+    wire = exchange(config, lambda _request: httpx.Response(200, stream=Stream()))
+    assert wire.split(b"\r\n\r\n", 1)[1] == chunk * chunks
+    assert max(retained, default=0) <= (1024 * 1024 if inspect else 0)
+    if inspect:
+        entries = [
+            json.loads(line)
+            for path in tmp_path.glob("*/requests.jsonl")
+            for line in path.read_text().splitlines()
+        ]
+        response = next(entry for entry in entries if entry["type"] == "response")
+        assert len(response["body"].encode()) == 1024 * 1024
+        assert response["body_truncated"] is True
+    else:
+        assert not list(tmp_path.glob("*/requests.jsonl"))