diff --git a/tests/tools/test_vision_tools.py b/tests/tools/test_vision_tools.py index 28e39c9657134..ed79cb7aaddff 100644 --- a/tests/tools/test_vision_tools.py +++ b/tests/tools/test_vision_tools.py @@ -179,10 +179,10 @@ class TestErrorLoggingExcInfo: from tools.vision_tools import _download_image with patch("tools.vision_tools.httpx.AsyncClient") as mock_client_cls: - mock_client = AsyncMock() + mock_client = MagicMock() mock_client.__aenter__ = AsyncMock(return_value=mock_client) mock_client.__aexit__ = AsyncMock(return_value=False) - mock_client.get = AsyncMock(side_effect=ConnectionError("network down")) + mock_client.stream.side_effect = ConnectionError("network down") mock_client_cls.return_value = mock_client dest = tmp_path / "image.jpg" @@ -365,29 +365,131 @@ class TestVisionSafetyGuards: } raise AssertionError(f"unexpected URL checked: {url}") - class FakeResponse: + class _FakeStreamResponse: url = "https://blocked.test/final.png" headers = {"content-length": "24"} - content = b"\x89PNG\r\n\x1a\n" + b"\x00" * 16 def raise_for_status(self): return None + async def aiter_bytes(self): + yield b"\x89PNG\r\n\x1a\n" + b"\x00" * 16 + + class _FakeAsyncStream: + def __init__(self, response): + self.response = response + + async def __aenter__(self): + return self.response + + async def __aexit__(self, exc_type, exc, tb): + return False + with ( patch("tools.vision_tools.check_website_access", side_effect=fake_check), patch("tools.vision_tools.httpx.AsyncClient") as mock_client_cls, pytest.raises(PermissionError, match="Blocked by website policy"), ): - mock_client = AsyncMock() + mock_client = MagicMock() mock_client.__aenter__ = AsyncMock(return_value=mock_client) mock_client.__aexit__ = AsyncMock(return_value=False) - mock_client.get = AsyncMock(return_value=FakeResponse()) + mock_client.stream.return_value = _FakeAsyncStream(_FakeStreamResponse()) mock_client_cls.return_value = mock_client - await _download_image("https://allowed.test/cat.png", tmp_path / "cat.png", max_retries=1) + await _download_image( + "https://allowed.test/cat.png", tmp_path / "cat.png", max_retries=1 + ) assert not (tmp_path / "cat.png").exists() + @pytest.mark.asyncio + async def test_download_enforces_size_cap_while_streaming(self, tmp_path): + """Streaming download rejects oversize payloads chunk-by-chunk and cleans up.""" + from tools.vision_tools import _download_image + + class _FakeStreamResponse: + url = "https://example.com/big.png" + headers = {} + + def raise_for_status(self): + return None + + async def aiter_bytes(self): + yield b"12345" + yield b"678901" + + class _FakeAsyncStream: + def __init__(self, response): + self.response = response + + async def __aenter__(self): + return self.response + + async def __aexit__(self, exc_type, exc, tb): + return False + + with ( + patch("tools.vision_tools._VISION_MAX_DOWNLOAD_BYTES", 10), + patch("tools.vision_tools.check_website_access", return_value=None), + patch("tools.vision_tools.httpx.AsyncClient") as mock_client_cls, + pytest.raises(ValueError, match="Image too large"), + ): + mock_client = MagicMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=False) + mock_client.stream.return_value = _FakeAsyncStream(_FakeStreamResponse()) + mock_client_cls.return_value = mock_client + + await _download_image( + "https://example.com/big.png", tmp_path / "big.png", max_retries=1 + ) + + assert not (tmp_path / "big.png").exists() + assert not list(tmp_path.glob(".big.png.*.tmp")) + + @pytest.mark.asyncio + async def test_download_ignores_malformed_content_length(self, tmp_path): + """Malformed Content-Length is ignored; streaming size cap still works.""" + from tools.vision_tools import _download_image + + body = b"\x89PNG\r\n\x1a\n" + b"\x00" * 16 + dest = tmp_path / "cat.png" + + class _FakeStreamResponse: + url = "https://example.com/cat.png" + headers = {"content-length": "not-a-number"} + + def raise_for_status(self): + return None + + async def aiter_bytes(self): + yield body[:4] + yield body[4:] + + class _FakeAsyncStream: + def __init__(self, response): + self.response = response + + async def __aenter__(self): + return self.response + + async def __aexit__(self, exc_type, exc, tb): + return False + + with ( + patch("tools.vision_tools.check_website_access", return_value=None), + patch("tools.vision_tools.httpx.AsyncClient") as mock_client_cls, + ): + mock_client = MagicMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=False) + mock_client.stream.return_value = _FakeAsyncStream(_FakeStreamResponse()) + mock_client_cls.return_value = mock_client + + await _download_image("https://example.com/cat.png", dest, max_retries=1) + + assert dest.read_bytes() == body + # --------------------------------------------------------------------------- # check_vision_requirements @@ -702,15 +804,18 @@ class TestDownloadRetryClassification: ) def _make_client_raising_status(self, status_code): - """AsyncClient whose response.raise_for_status() raises HTTPStatusError.""" + """AsyncClient whose stream response.raise_for_status() raises HTTPStatusError.""" mock_response = MagicMock() mock_response.raise_for_status = MagicMock( side_effect=self._status_error(status_code) ) - mock_client = AsyncMock() + mock_stream = MagicMock() + mock_stream.__aenter__ = AsyncMock(return_value=mock_response) + mock_stream.__aexit__ = AsyncMock(return_value=False) + mock_client = MagicMock() mock_client.__aenter__ = AsyncMock(return_value=mock_client) mock_client.__aexit__ = AsyncMock(return_value=False) - mock_client.get = AsyncMock(return_value=mock_response) + mock_client.stream = MagicMock(return_value=mock_stream) return mock_client def test_is_retryable_classification(self): @@ -746,7 +851,7 @@ class TestDownloadRetryClassification: "https://example.com/flaky.jpg", tmp_path / "y.jpg", max_retries=3 ) # All three attempts used, two backoff sleeps between them. - assert mock_client.get.await_count == 3 + assert mock_client.stream.call_count == 3 assert mock_sleep.await_count == 2 diff --git a/tools/vision_tools.py b/tools/vision_tools.py index cef07ee023733..5e4de78eb2983 100644 --- a/tools/vision_tools.py +++ b/tools/vision_tools.py @@ -403,6 +403,77 @@ def _is_retryable_download_error(error: Exception) -> bool: return True +async def _stream_download_to_file( + client, + url: str, + destination: Path, + max_bytes: int, + *, + headers: dict, + media_label: str = "Image", +) -> Path: + """Stream an HTTP download to *destination* via a temp file with a running size cap. + + Uses ``client.stream("GET", ...)`` so the response body is never fully + buffered in memory — chunks are written to a temp file and the running + byte count is checked against *max_bytes* after each chunk. On success + the temp file is atomically replaced onto *destination*; on failure the + temp file is deleted. + + A ``Content-Length`` header, when present and parseable, is used for an + early rejection before any bytes are streamed, but the streaming cap is + the authoritative guard (servers can omit or lie about the header). + """ + from utils import atomic_replace + + async with client.stream("GET", url, headers=headers) as response: + response.raise_for_status() + + # Early rejection via Content-Length when present and valid. + cl = response.headers.get("content-length") + if cl: + try: + declared_size = int(cl) + except ValueError: + declared_size = None + if declared_size is not None and declared_size > max_bytes: + raise ValueError( + f"{media_label} too large ({declared_size} bytes, max {max_bytes})" + ) + + final_url = str(response.url) + blocked = check_website_access(final_url) + if blocked: + raise PermissionError(blocked["message"]) + + tmp_destination = destination.with_name( + f".{destination.name}.{uuid.uuid4().hex}.tmp" + ) + bytes_written = 0 + try: + with tmp_destination.open("wb") as f: + async for chunk in response.aiter_bytes(): + if not chunk: + continue + bytes_written += len(chunk) + if bytes_written > max_bytes: + raise ValueError( + f"{media_label} too large ({bytes_written} bytes, max {max_bytes})" + ) + f.write(chunk) + atomic_replace(tmp_destination, destination) + except Exception: + try: + tmp_destination.unlink(missing_ok=True) + except OSError: + logger.debug( + "Could not delete partial download: %s", tmp_destination, exc_info=True + ) + raise + + return destination + + async def _download_image(image_url: str, destination: Path, max_retries: int = 3) -> Path: """ Download an image from a URL to a local destination (async) with retry logic. @@ -447,43 +518,28 @@ async def _download_image(image_url: str, destination: Path, max_retries: int = from tools.url_safety import create_ssrf_safe_async_client - # Download the image with appropriate headers using async httpx - # Enable follow_redirects to handle image CDNs that redirect (e.g., Imgur, Picsum) + # Download the image with appropriate headers using async httpx. + # Enable follow_redirects to handle image CDNs that redirect (e.g., Imgur, Picsum). # SSRF: the client validates DNS at TCP connect time; event_hooks # validate each redirect target against private IP ranges. + # Streaming: body is written chunk-by-chunk to a temp file so the + # size cap bounds memory, not just disk. async with create_ssrf_safe_async_client( timeout=_VISION_DOWNLOAD_TIMEOUT, follow_redirects=True, event_hooks={"response": [_ssrf_redirect_guard]}, ) as client: - response = await client.get( + await _stream_download_to_file( + client, image_url, + destination, + _VISION_MAX_DOWNLOAD_BYTES, headers={ "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36", "Accept": "image/*,*/*;q=0.8", }, + media_label="Image", ) - response.raise_for_status() - - # Reject overly large images early via Content-Length header. - cl = response.headers.get("content-length") - if cl and int(cl) > _VISION_MAX_DOWNLOAD_BYTES: - raise ValueError( - f"Image too large ({int(cl)} bytes, max {_VISION_MAX_DOWNLOAD_BYTES})" - ) - - final_url = str(response.url) - blocked = check_website_access(final_url) - if blocked: - raise PermissionError(blocked["message"]) - - # Save the image content (double-check actual size) - body = response.content - if len(body) > _VISION_MAX_DOWNLOAD_BYTES: - raise ValueError( - f"Image too large ({len(body)} bytes, max {_VISION_MAX_DOWNLOAD_BYTES})" - ) - destination.write_bytes(body) return destination except Exception as e: @@ -1607,32 +1663,17 @@ async def _download_video(video_url: str, destination: Path, max_retries: int = follow_redirects=True, event_hooks={"response": [_ssrf_redirect_guard]}, ) as client: - response = await client.get( + await _stream_download_to_file( + client, video_url, + destination, + _MAX_VIDEO_BASE64_BYTES, headers={ "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36", "Accept": "video/*,*/*;q=0.8", }, + media_label="Video", ) - response.raise_for_status() - - cl = response.headers.get("content-length") - if cl and int(cl) > _MAX_VIDEO_BASE64_BYTES: - raise ValueError( - f"Video too large ({int(cl)} bytes, max {_MAX_VIDEO_BASE64_BYTES})" - ) - - final_url = str(response.url) - blocked = check_website_access(final_url) - if blocked: - raise PermissionError(blocked["message"]) - - body = response.content - if len(body) > _MAX_VIDEO_BASE64_BYTES: - raise ValueError( - f"Video too large ({len(body)} bytes, max {_MAX_VIDEO_BASE64_BYTES})" - ) - destination.write_bytes(body) return destination except Exception as e: