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()