From fba2be9f2bcb3527d7c4680e6a566feca2c29836 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 3 Jul 2026 00:31:50 +0000 Subject: [PATCH] Send SSE keepalives during image generation so idle-timeout proxies don't kill the stream gpt-image-2 generations regularly take 60-120s, and the streaming responder ran the whole provider call before yielding a single byte. With nothing on the wire, an idle-timeout proxy between the enclave and the browser resets the HTTP/2 stream at ~60s, surfacing client-side as ERR_HTTP2_PROTOCOL_ERROR. Run the provider call on a worker thread and yield SSE comment frames (": keepalive") every 10s while it executes. The frames are sealed as ordinary chunked-OHTTP chunks, so bytes keep flowing across every hop (enclave -> relay -> browser), and SSE parsers ignore comment lines so the client contract is unchanged. Applied to both long-silent paths: the endpoint-based image models (gpt-image-2, grok-2-image, Seedream, GLM-Image) and the Gemini inline-image invoke in the chat controller. Also give the provider image call its own 300s read timeout instead of the shared chat client default, since image generation is far slower than chat. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_017bwXiKz9YQsuhTyRcwEnKz --- tee_gateway/controllers/chat_controller.py | 10 ++- tee_gateway/image_generation.py | 71 +++++++++++++++- tee_gateway/test/test_image_generation.py | 99 ++++++++++++++++++++++ 3 files changed, 174 insertions(+), 6 deletions(-) diff --git a/tee_gateway/controllers/chat_controller.py b/tee_gateway/controllers/chat_controller.py index 1392b6a..4a3ef2c 100644 --- a/tee_gateway/controllers/chat_controller.py +++ b/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 ) diff --git a/tee_gateway/image_generation.py b/tee_gateway/image_generation.py index 07776fa..7594fb6 100644 --- a/tee_gateway/image_generation.py +++ b/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"}], diff --git a/tee_gateway/test/test_image_generation.py b/tee_gateway/test/test_image_generation.py index c6c9c6b..43c45e9 100644 --- a/tee_gateway/test/test_image_generation.py +++ b/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()