Add SSE keepalives during long image generation calls

tee_gateway/controllers/chat_controller.py

@@ -38,6 +38,7 @@
 from tee_gateway.image_generation import (
     create_image_generation_response,
     create_image_generation_streaming_response,
+    run_with_sse_keepalives,
 )
 from tee_gateway.model_registry import get_model_config
 from tee_gateway.pricing import compute_session_cost
@@ -484,8 +485,13 @@ def generate():
                 if image_output_model:
                     # Image generation isn't a token stream: invoke once, emit
                     # any caption text as a content delta, and carry the image
-                    # out-of-band on the final frame (it is not signed).
-                    response = model.invoke(langchain_messages)
+                    # out-of-band on the final frame (it is not signed). The
+                    # invoke runs on a worker thread with SSE keepalives so the
+                    # long silent provider call doesn't trip idle-timeout
+                    # proxies between the enclave and the browser.
+                    response = yield from run_with_sse_keepalives(
+                        lambda: model.invoke(langchain_messages)
+                    )
                     text_content, generated_images = _split_text_and_images(
                         response.content
                     )

tee_gateway/image_generation.py

@@ -26,9 +26,10 @@
 import ipaddress
 import json
 import logging
+import threading
 import time
 import uuid
-from typing import Any, List, Optional
+from typing import Any, Callable, Generator, List, Optional, TypeVar
 from urllib.parse import urlparse
 
 import httpx
@@ -59,6 +60,55 @@
     "zai": "zai_http_client",
 }
 
+# Image generation is far slower than chat (gpt-image regularly takes 1-2
+# minutes), so the provider call gets its own generous read timeout instead of
+# the shared chat client's default.
+_IMAGE_GENERATION_TIMEOUT = httpx.Timeout(timeout=300.0, connect=15.0)
+
+# SSE comment frame emitted while the provider call runs. Comment lines are
+# ignored by SSE parsers (the app's stream parser only reads ``data:`` lines)
+# but keep bytes flowing on the wire: without them the connection is silent for
+# the entire generation, and idle-timeout proxies between the enclave and the
+# browser reset the stream at ~60s (surfacing client-side as
+# ERR_HTTP2_PROTOCOL_ERROR).
+_SSE_KEEPALIVE_FRAME = ": keepalive\n\n"
+_KEEPALIVE_INTERVAL_SECONDS = 10.0
+
+_T = TypeVar("_T")
+
+
+def run_with_sse_keepalives(fn: Callable[[], _T]) -> Generator[str, None, _T]:
+    """Run ``fn`` on a worker thread, yielding SSE keepalive comments while it runs.
+
+    Use with ``yield from`` inside a streaming response generator: the caller
+    receives ``fn``'s return value, and any exception ``fn`` raised is re-raised
+    at the call site. The OHTTP encrypter downstream holds each chunk until the
+    next arrives (one-chunk look-ahead for the final marker), so the wire sees
+    bytes at most every 2x the interval — still far under the ~60s idle
+    timeouts this defends against. If the client disconnects mid-generation the
+    generator is closed and the daemon worker is left to finish on its own (the
+    provider call is not cancellable).
+    """
+    outcome: dict[str, Any] = {}
+
+    def _work() -> None:
+        try:
+            outcome["result"] = fn()
+        except BaseException as exc:  # re-raised on the streaming thread below
+            outcome["error"] = exc
+
+    worker = threading.Thread(target=_work, name="sse-keepalive-call", daemon=True)
+    worker.start()
+    while True:
+        worker.join(timeout=_KEEPALIVE_INTERVAL_SECONDS)
+        if not worker.is_alive():
+            break
+        yield _SSE_KEEPALIVE_FRAME
+    if "error" in outcome:
+        raise outcome["error"]
+    return outcome["result"]  # type: ignore[no-any-return]
+
+
 # Bounds on the URL fetch (egress hardening). Provider images are well under the
 # size cap; the redirect cap stops a redirect chain from being chased off-host.
 _MAX_IMAGE_BYTES = 25 * 1024 * 1024  # 25 MiB
@@ -294,7 +344,12 @@ def generate_images(
             form["response_format"] = cfg.image_response_format
         if cfg.image_extra_params:
             form.update({k: str(v) for k, v in cfg.image_extra_params.items()})
-        resp = client.post(edit_endpoint, data=form, files=uploads)
+        resp = client.post(
+            edit_endpoint,
+            data=form,
+            files=uploads,
+            timeout=_IMAGE_GENERATION_TIMEOUT,
+        )
     else:
         # No edit endpoint (or nothing uploadable): JSON generations. Inline
         # references only ride along for providers that carry them there and
@@ -303,6 +358,7 @@ def generate_images(
         resp = client.post(
             _IMAGE_GENERATION_PATH,
             json=_build_generations_payload(cfg, prompt, count, json_refs),
+            timeout=_IMAGE_GENERATION_TIMEOUT,
         )
     resp.raise_for_status()
     data = resp.json().get("data", []) or []
@@ -469,11 +525,18 @@ def create_image_generation_streaming_response(
     chat_request: CreateChatCompletionRequest, request_bytes: bytes
 ):
     """Streaming image generation: image gen is not a token stream, so we invoke
-    once and emit the result on the final SSE frame (mirrors the Gemini path)."""
+    once and emit the result on the final SSE frame (mirrors the Gemini path).
+
+    The provider call runs on a worker thread with SSE keepalive comments
+    yielded while it executes — image generation is silent for 60-120s, and
+    without bytes on the wire idle-timeout proxies between the enclave and the
+    browser reset the stream at ~60s."""
 
     def generate():
         try:
-            result = _run_image_generation(chat_request, request_bytes)
+            result = yield from run_with_sse_keepalives(
+                lambda: _run_image_generation(chat_request, request_bytes)
+            )
 
             final_data: dict[str, Any] = {
                 "choices": [{"delta": {}, "index": 0, "finish_reason": "stop"}],

tee_gateway/test/test_image_generation.py

@@ -13,6 +13,7 @@
 fetch is patched, and a stub price feed is injected.
 """
 
+import time
 import unittest
 from decimal import Decimal
 from unittest.mock import MagicMock, patch
@@ -536,5 +537,103 @@ def test_zero_images_is_free(self):
         self.assertEqual(cost.cost_opg, 0)
 
 
+class TestStreamingKeepalives(unittest.TestCase):
+    """SSE keepalives during the (long, silent) provider image call.
+
+    Image generation takes 60-120s with no bytes on the wire; idle-timeout
+    proxies between the enclave and the browser reset the stream at ~60s
+    unless keepalive comment frames keep it warm. These tests pin that the
+    streaming responder emits keepalives while the generation runs and that
+    the final signed frame / error frame semantics are unchanged.
+    """
+
+    _RESULT = {
+        "images": ["data:image/png;base64,aGVsbG8="],
+        "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
+        "opengradient": None,
+        "tee_signature": "sig",
+        "tee_request_hash": "req",
+        "tee_output_hash": "out",
+        "tee_timestamp": 1,
+        "tee_id": "0xabc",
+    }
+
+    @staticmethod
+    def _chat_request() -> MagicMock:
+        request = MagicMock()
+        request.model = GPT_IMAGE
+        return request
+
+    def _collect_frames(self, run_side_effect) -> list[str]:
+        with (
+            patch.object(image_generation, "_KEEPALIVE_INTERVAL_SECONDS", 0.01),
+            patch.object(
+                image_generation,
+                "_run_image_generation",
+                side_effect=run_side_effect,
+            ),
+        ):
+            response = image_generation.create_image_generation_streaming_response(
+                self._chat_request(), b"{}"
+            )
+            return [
+                chunk.decode() if isinstance(chunk, bytes) else chunk
+                for chunk in response.response
+            ]
+
+    def test_keepalives_flow_while_generation_runs(self):
+        def slow_generation(*_args, **_kwargs):
+            time.sleep(0.05)
+            return dict(self._RESULT)
+
+        frames = self._collect_frames(slow_generation)
+
+        keepalives = [f for f in frames if f == ": keepalive\n\n"]
+        self.assertGreater(len(keepalives), 0)
+        # Keepalives come strictly before the final data frames.
+        first_data = next(i for i, f in enumerate(frames) if f.startswith("data:"))
+        self.assertTrue(
+            all(f == ": keepalive\n\n" for f in frames[:first_data]),
+        )
+        # Final signed frame and [DONE] are intact.
+        self.assertIn("tee_signature", frames[first_data])
+        self.assertIn("data:image/png;base64,aGVsbG8=", frames[first_data])
+        self.assertEqual(frames[-1], "data: [DONE]\n\n")
+
+    def test_generation_error_still_yields_error_frame(self):
+        def failing_generation(*_args, **_kwargs):
+            time.sleep(0.05)
+            raise RuntimeError("provider exploded")
+
+        frames = self._collect_frames(failing_generation)
+
+        self.assertGreater(len([f for f in frames if f == ": keepalive\n\n"]), 0)
+        self.assertIn("provider exploded", frames[-1])
+        self.assertIn("RuntimeError", frames[-1])
+
+    def test_fast_generation_emits_no_keepalives(self):
+        frames = self._collect_frames(lambda *_a, **_k: dict(self._RESULT))
+        self.assertNotIn(": keepalive\n\n", frames)
+
+    def test_run_with_sse_keepalives_returns_value_and_raises(self):
+        def drive(gen):
+            frames = []
+            try:
+                while True:
+                    frames.append(next(gen))
+            except StopIteration as stop:
+                return frames, stop.value
+
+        with patch.object(image_generation, "_KEEPALIVE_INTERVAL_SECONDS", 0.01):
+            _, value = drive(image_generation.run_with_sse_keepalives(lambda: 42))
+            self.assertEqual(value, 42)
+
+            def boom():
+                raise ValueError("bad")
+
+            with self.assertRaises(ValueError):
+                drive(image_generation.run_with_sse_keepalives(boom))
+
+
 if __name__ == "__main__":
     unittest.main()