diff --git a/agent/image_gen_provider.py b/agent/image_gen_provider.py index 47f65c1b343..a7f1b8c31ff 100644 --- a/agent/image_gen_provider.py +++ b/agent/image_gen_provider.py @@ -191,6 +191,88 @@ def save_b64_image( return path +# Extension inference for save_url_image — keep small and explicit. We don't +# want to import mimetypes for a handful of formats every image_gen provider +# actually returns, and we never want to inherit a content-type that points +# at HTML or JSON when the API gives us a degenerate response. +_URL_IMAGE_CONTENT_TYPES = { + "image/png": "png", + "image/jpeg": "jpg", + "image/jpg": "jpg", + "image/webp": "webp", + "image/gif": "gif", +} + + +def save_url_image( + url: str, + *, + prefix: str = "image", + timeout: float = 60.0, + max_bytes: int = 25 * 1024 * 1024, +) -> Path: + """Download an image URL and write it under ``$HERMES_HOME/cache/images/``. + + Used by providers (xAI, fallback OpenAI) whose API returns an *ephemeral* + URL instead of inline base64 — those URLs frequently expire before a + downstream consumer (Telegram ``send_photo``, browser fetch) can resolve + them, so we materialise the bytes locally at tool-completion time. + Mirrors :func:`save_b64_image`'s shape so providers can swap in one line. + + Returns the absolute :class:`Path` to the saved file. Raises on any + network / HTTP / oversize / non-image-content-type error so callers can + fall back to returning the bare URL with a clear error message. + """ + import requests + + response = requests.get(url, timeout=timeout, stream=True) + response.raise_for_status() + + # Infer extension from the response content-type, falling back to the + # URL suffix when xAI / OpenAI omit a precise type (some CDNs return + # ``application/octet-stream``). Defaults to ``png``. + content_type = (response.headers.get("Content-Type") or "").split(";", 1)[0].strip().lower() + extension = _URL_IMAGE_CONTENT_TYPES.get(content_type) + if extension is None: + url_path = url.split("?", 1)[0].lower() + for ext in ("png", "jpg", "jpeg", "webp", "gif"): + if url_path.endswith(f".{ext}"): + extension = "jpg" if ext == "jpeg" else ext + break + if extension is None: + extension = "png" + + ts = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") + short = uuid.uuid4().hex[:8] + path = _images_cache_dir() / f"{prefix}_{ts}_{short}.{extension}" + + bytes_written = 0 + with path.open("wb") as fh: + for chunk in response.iter_content(chunk_size=64 * 1024): + if not chunk: + continue + bytes_written += len(chunk) + if bytes_written > max_bytes: + fh.close() + try: + path.unlink() + except OSError: + pass + raise ValueError( + f"Image at {url} exceeds {max_bytes // (1024 * 1024)}MB cap; refusing to cache." + ) + fh.write(chunk) + + if bytes_written == 0: + try: + path.unlink() + except OSError: + pass + raise ValueError(f"Image at {url} returned 0 bytes; refusing to cache.") + + return path + + def success_response( *, image: str, diff --git a/plugins/image_gen/openai/__init__.py b/plugins/image_gen/openai/__init__.py index c1a719f9102..448f5bc45af 100644 --- a/plugins/image_gen/openai/__init__.py +++ b/plugins/image_gen/openai/__init__.py @@ -33,6 +33,7 @@ from agent.image_gen_provider import ( error_response, resolve_aspect_ratio, save_b64_image, + save_url_image, success_response, ) @@ -266,9 +267,21 @@ class OpenAIImageGenProvider(ImageGenProvider): ) image_ref = str(saved_path) elif url: - # Defensive — gpt-image-2 returns b64 today, but fall back - # gracefully if the API ever changes. - image_ref = url + # Defensive — gpt-image-2 returns b64 today, but OpenAI's API + # has previously returned URLs. Cache the bytes locally so the + # gateway never tries to fetch an ephemeral / signed URL after + # it expires — same rationale as the xAI provider (#26942). + try: + saved_path = save_url_image(url, prefix=f"openai_{tier_id}") + except Exception as exc: + logger.warning( + "OpenAI image URL %s could not be cached (%s); falling back to bare URL.", + url, + exc, + ) + image_ref = url + else: + image_ref = str(saved_path) else: return error_response( error="OpenAI response contained neither b64_json nor URL", diff --git a/plugins/image_gen/xai/__init__.py b/plugins/image_gen/xai/__init__.py index d5aac4eccdd..a8982393f7e 100644 --- a/plugins/image_gen/xai/__init__.py +++ b/plugins/image_gen/xai/__init__.py @@ -29,6 +29,7 @@ from agent.image_gen_provider import ( error_response, resolve_aspect_ratio, save_b64_image, + save_url_image, success_response, ) from tools.xai_http import hermes_xai_user_agent, resolve_xai_http_credentials @@ -281,7 +282,24 @@ class XAIImageGenProvider(ImageGenProvider): ) image_ref = str(saved_path) elif url: - image_ref = url + # xAI's grok-imagine-image returns ephemeral ``imgen.x.ai/xai-tmp-*`` + # URLs that 404 within minutes — by the time Telegram's + # ``send_photo`` or any downstream consumer fetches them, the + # asset is gone (#26942). Materialise the bytes locally at + # tool-completion time so the gateway has a stable file path to + # upload, mirroring the b64 branch above and the audio_cache + # pattern used by text_to_speech. + try: + saved_path = save_url_image(url, prefix=f"xai_{model_id}") + except Exception as exc: + logger.warning( + "xAI image URL %s could not be cached (%s); falling back to bare URL.", + url, + exc, + ) + image_ref = url + else: + image_ref = str(saved_path) else: return error_response( error="xAI response contained neither b64_json nor URL", diff --git a/tests/agent/test_save_url_image.py b/tests/agent/test_save_url_image.py new file mode 100644 index 00000000000..6a63413f74e --- /dev/null +++ b/tests/agent/test_save_url_image.py @@ -0,0 +1,168 @@ +"""Direct tests for ``agent.image_gen_provider.save_url_image`` (#26942). + +These exercise the helper against a real in-process HTTP server — no +``requests.get`` mocking — so we catch the kinds of issues a mocked +unit test won't: content-type parsing, partial-write cleanup, the +oversize cap, the empty-body refusal, and the cache directory it +actually writes to. + +Pre-fix the helper didn't exist; xAI URL responses were returned bare +and the gateway 404'd at ``send_photo`` time. +""" + +from __future__ import annotations + +import http.server +import socketserver +import threading + +import pytest + + +PNG_1PX = bytes.fromhex( + "89504e470d0a1a0a0000000d49484452000000010000000108020000009077" + "53de00000010494441547801635c0e000000feff03000006000557bfabd400" + "00000049454e44ae426082" +) + + +class _TinyImageHandler(http.server.BaseHTTPRequestHandler): + """Tiny HTTP server that mimics the shapes save_url_image must handle.""" + + def do_GET(self): # noqa: N802 + if self.path == "/image.png": + self.send_response(200) + self.send_header("Content-Type", "image/png") + self.send_header("Content-Length", str(len(PNG_1PX))) + self.end_headers() + self.wfile.write(PNG_1PX) + elif self.path == "/image.jpg": + self.send_response(200) + self.send_header("Content-Type", "image/jpeg") + self.end_headers() + self.wfile.write(PNG_1PX) # bytes don't have to be a real jpeg + elif self.path == "/oversize": + self.send_response(200) + self.send_header("Content-Type", "image/png") + self.end_headers() + chunk = b"\x00" * 65536 + for _ in range(64): # 4 MiB + self.wfile.write(chunk) + elif self.path == "/empty": + self.send_response(200) + self.send_header("Content-Type", "image/png") + self.send_header("Content-Length", "0") + self.end_headers() + elif self.path == "/404": + self.send_response(404) + self.end_headers() + elif self.path == "/no-type-with-url-ext.jpg": + self.send_response(200) + self.send_header("Content-Type", "application/octet-stream") + self.end_headers() + self.wfile.write(PNG_1PX) + elif self.path == "/no-type-no-ext": + self.send_response(200) + self.end_headers() + self.wfile.write(PNG_1PX) + else: + self.send_response(404) + self.end_headers() + + def log_message(self, *args, **kw): # noqa: D401 + return + + +@pytest.fixture +def http_server(tmp_path, monkeypatch): + """Spin up a localhost HTTP server and isolate HERMES_HOME under tmp_path.""" + monkeypatch.setenv("HERMES_HOME", str(tmp_path / ".hermes")) + (tmp_path / ".hermes").mkdir() + + # Force the constants/image cache helpers to re-read HERMES_HOME. + import sys + for mod in list(sys.modules): + if mod.startswith("hermes_constants") or mod.startswith("agent.image_gen_provider"): + sys.modules.pop(mod, None) + + httpd = socketserver.TCPServer(("127.0.0.1", 0), _TinyImageHandler) + port = httpd.server_address[1] + thread = threading.Thread(target=httpd.serve_forever, daemon=True) + thread.start() + yield f"http://127.0.0.1:{port}", httpd + httpd.shutdown() + + +class TestSaveUrlImage: + def test_writes_real_bytes_to_hermes_home_cache(self, http_server): + base, _ = http_server + from agent.image_gen_provider import save_url_image + + path = save_url_image(f"{base}/image.png", prefix="xai_test") + + assert path.exists() + assert path.read_bytes() == PNG_1PX + # The cache directory must be under HERMES_HOME — gateway cleanup + # relies on this being the canonical location. + assert "cache/images" in str(path) + assert path.suffix == ".png" + + def test_extension_inferred_from_content_type(self, http_server): + base, _ = http_server + from agent.image_gen_provider import save_url_image + + path = save_url_image(f"{base}/image.jpg", prefix="xai_test") + assert path.suffix == ".jpg", "image/jpeg → .jpg" + + def test_extension_falls_back_to_url_suffix(self, http_server): + """Some CDNs send ``application/octet-stream`` — the URL suffix wins then.""" + base, _ = http_server + from agent.image_gen_provider import save_url_image + + path = save_url_image(f"{base}/no-type-with-url-ext.jpg", prefix="xai_test") + assert path.suffix == ".jpg" + + def test_extension_defaults_to_png_when_unknowable(self, http_server): + base, _ = http_server + from agent.image_gen_provider import save_url_image + + path = save_url_image(f"{base}/no-type-no-ext", prefix="xai_test") + assert path.suffix == ".png" + + def test_404_raises(self, http_server): + """HTTP errors must propagate — caller decides whether to fall back.""" + base, _ = http_server + from agent.image_gen_provider import save_url_image + import requests as req_lib + + with pytest.raises(req_lib.HTTPError): + save_url_image(f"{base}/404") + + def test_empty_body_raises_without_writing_file(self, http_server): + """0-byte responses are not images — refuse to cache.""" + base, _ = http_server + from agent.image_gen_provider import save_url_image + + with pytest.raises(ValueError, match="0 bytes"): + save_url_image(f"{base}/empty") + + def test_oversize_raises_and_cleans_up(self, http_server, tmp_path): + """Oversize downloads must NOT leak a partial file into the cache.""" + base, _ = http_server + from agent.image_gen_provider import save_url_image, _images_cache_dir + + cache_dir = _images_cache_dir() + before = set(cache_dir.glob("*")) + with pytest.raises(ValueError, match="exceeds"): + save_url_image(f"{base}/oversize", max_bytes=1024 * 1024) + after = set(cache_dir.glob("*")) + assert after == before, "partial file leaked into cache after oversize cap" + + def test_unique_filenames_avoid_collision(self, http_server): + """Two back-to-back saves of the same URL must produce different paths.""" + base, _ = http_server + from agent.image_gen_provider import save_url_image + + path1 = save_url_image(f"{base}/image.png", prefix="xai_collision") + path2 = save_url_image(f"{base}/image.png", prefix="xai_collision") + assert path1 != path2, "filename collision — uuid suffix isn't doing its job" diff --git a/tests/plugins/image_gen/test_openai_provider.py b/tests/plugins/image_gen/test_openai_provider.py index 670722efbde..6411996130e 100644 --- a/tests/plugins/image_gen/test_openai_provider.py +++ b/tests/plugins/image_gen/test_openai_provider.py @@ -229,14 +229,43 @@ class TestGenerate: assert result["success"] is False assert result["error_type"] == "empty_response" - def test_url_fallback_if_api_changes(self, provider): - """Defensive: if OpenAI ever returns URL instead of b64, pass through.""" + def test_url_response_is_cached_locally(self, provider): + """OpenAI URL response (if API ever returns one) is cached locally. + + Pre-fix this asserted the bare URL passed through; symmetric to the + xAI #26942 fix. Even though gpt-image-2 returns b64 today, every + ``image_gen`` provider must guarantee the gateway gets a stable + file path so ephemeral signed URLs can't expire mid-flight. + """ fake_client = MagicMock() fake_client.images.generate.return_value = _fake_response( b64=None, url="https://example.com/img.png", ) - with _patched_openai(fake_client): + with _patched_openai(fake_client), patch( + "plugins.image_gen.openai.save_url_image", + return_value=Path("/tmp/openai_gpt-image-2_20260524_000000_deadbeef.png"), + ) as mock_save_url: + result = provider.generate("a cat") + + assert result["success"] is True + assert result["image"].startswith("/") + assert "example.com" not in result["image"] + mock_save_url.assert_called_once() + + def test_url_response_falls_back_to_bare_url_when_download_fails(self, provider): + """Cache failure must not turn into a tool error — symmetric with xAI.""" + import requests as req_lib + + fake_client = MagicMock() + fake_client.images.generate.return_value = _fake_response( + b64=None, url="https://example.com/img.png", + ) + + with _patched_openai(fake_client), patch( + "plugins.image_gen.openai.save_url_image", + side_effect=req_lib.HTTPError("404 from CDN"), + ): result = provider.generate("a cat") assert result["success"] is True diff --git a/tests/plugins/image_gen/test_xai_provider.py b/tests/plugins/image_gen/test_xai_provider.py index 88ce31813e4..f921fe2e291 100644 --- a/tests/plugins/image_gen/test_xai_provider.py +++ b/tests/plugins/image_gen/test_xai_provider.py @@ -5,6 +5,7 @@ from __future__ import annotations import json import os +from pathlib import Path from unittest.mock import MagicMock, patch import pytest @@ -142,21 +143,75 @@ class TestGenerate: assert result["model"] == "grok-imagine-image" def test_successful_url_response(self): + """xAI URL response is cached locally — #26942 contract. + + Pre-fix this asserted ``result["image"] == ""``, which + was exactly the bug: xAI's ``imgen.x.ai/xai-tmp-*`` URLs expire fast + and the gateway 404'd by ``send_photo`` time. Post-fix the URL + bytes are downloaded at tool-completion and the result carries an + absolute filesystem path the gateway can upload from. + """ from plugins.image_gen.xai import XAIImageGenProvider mock_resp = MagicMock() mock_resp.status_code = 200 mock_resp.raise_for_status = MagicMock() mock_resp.json.return_value = { - "data": [{"url": "https://xai.image/result.png"}], + "data": [{"url": "https://imgen.x.ai/xai-tmp-imgen-test.jpeg"}], } - with patch("plugins.image_gen.xai.requests.post", return_value=mock_resp): + with patch("plugins.image_gen.xai.requests.post", return_value=mock_resp), \ + patch( + "plugins.image_gen.xai.save_url_image", + return_value=Path("/tmp/xai_grok-imagine-image_20260524_000000_deadbeef.jpg"), + ) as mock_save_url: provider = XAIImageGenProvider() result = provider.generate(prompt="A cat playing piano") assert result["success"] is True - assert result["image"] == "https://xai.image/result.png" + assert result["image"].startswith("/"), ( + f"URL response must be cached to an absolute path, got {result['image']!r}" + ) + assert "imgen.x.ai" not in result["image"], ( + "ephemeral xAI URL must not leak into result.image — caller will 404" + ) + # The downloader should have been called exactly once with the URL + # and an xai-prefixed cache filename. + mock_save_url.assert_called_once() + call_args, call_kwargs = mock_save_url.call_args + assert call_args[0] == "https://imgen.x.ai/xai-tmp-imgen-test.jpeg" + assert call_kwargs.get("prefix", "").startswith("xai_") + + def test_url_response_falls_back_to_bare_url_when_download_fails(self): + """If caching the URL fails (network blip, 404 in-flight), the + provider must NOT hard-error — fall through to returning the bare + URL so the agent surface at least sees *something*. The gateway's + existing URL-send fallback then has a chance to succeed; if it + too 404s, the user gets the original (now legible) error rather + than an opaque "image generation failed" tool result. + """ + import requests as req_lib + from plugins.image_gen.xai import XAIImageGenProvider + + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.raise_for_status = MagicMock() + mock_resp.json.return_value = { + "data": [{"url": "https://imgen.x.ai/xai-tmp-imgen-already-404.jpeg"}], + } + + with patch("plugins.image_gen.xai.requests.post", return_value=mock_resp), \ + patch( + "plugins.image_gen.xai.save_url_image", + side_effect=req_lib.HTTPError("404 from CDN"), + ): + provider = XAIImageGenProvider() + result = provider.generate(prompt="A cat playing piano") + + assert result["success"] is True, ( + "Cache failure must not turn into a tool error — gateway gets a chance to retry" + ) + assert result["image"] == "https://imgen.x.ai/xai-tmp-imgen-already-404.jpeg" def test_api_error(self): import requests as req_lib