diff --git a/astrbot/core/provider/sources/openai_source.py b/astrbot/core/provider/sources/openai_source.py index 15fd6b72f..7d5a9ffd9 100644 --- a/astrbot/core/provider/sources/openai_source.py +++ b/astrbot/core/provider/sources/openai_source.py @@ -371,6 +371,9 @@ class ProviderOpenAIOfficial(Provider): default_headers=self.custom_headers, base_url=provider_config.get("api_base") or None, timeout=self.timeout, + # Retry is handled by retry_provider_request(); disable the + # SDK built-in retry to avoid stacking request attempts. + max_retries=0, http_client=self._create_http_client(provider_config), ) else: @@ -380,6 +383,9 @@ class ProviderOpenAIOfficial(Provider): base_url=provider_config.get("api_base") or None, default_headers=self.custom_headers, timeout=self.timeout, + # Retry is handled by retry_provider_request(); disable the + # SDK built-in retry to avoid stacking request attempts. + max_retries=0, http_client=self._create_http_client(provider_config), ) diff --git a/astrbot/core/utils/media_utils.py b/astrbot/core/utils/media_utils.py index 966cc4219..1f0c445b7 100644 --- a/astrbot/core/utils/media_utils.py +++ b/astrbot/core/utils/media_utils.py @@ -42,6 +42,11 @@ IMAGE_COMPRESS_DEFAULT_MAX_SIZE = 1280 IMAGE_COMPRESS_DEFAULT_QUALITY = 95 IMAGE_COMPRESS_DEFAULT_OPTIMIZE = True IMAGE_COMPRESS_DEFAULT_MIN_FILE_SIZE_MB = 1.0 +# Model image inputs larger than this are skipped before decoding. +MODEL_IMAGE_MAX_INPUT_BYTES = 32 * 1024 * 1024 +# Original encoded bytes are reused only for small stills; larger inputs are +# re-encoded so the output stays bounded by pixel size and quality. +MODEL_IMAGE_REUSE_MAX_BYTES = 2 * 1024 * 1024 MEDIA_MIME_EXTENSIONS = { "audio/wav": ".wav", @@ -67,15 +72,17 @@ MEDIA_MIME_EXTENSIONS = { "video/quicktime": ".mov", } -IMAGE_FORMAT_MIME_TYPES = { - "JPEG": "image/jpeg", - "PNG": "image/png", - "GIF": "image/gif", - "WEBP": "image/webp", - "BMP": "image/bmp", - "TIFF": "image/tiff", - "AVIF": "image/avif", -} +# Magic-byte prefixes for O(1) image MIME sniffing; unknown headers fall +# back to the caller-provided default instead of decoding the file. +_IMAGE_MAGIC_MIME_TYPES = ( + (b"\x89PNG\r\n\x1a\n", "image/png"), + (b"\xff\xd8\xff", "image/jpeg"), + (b"GIF87a", "image/gif"), + (b"GIF89a", "image/gif"), + (b"BM", "image/bmp"), + (b"II*\x00", "image/tiff"), + (b"MM\x00*", "image/tiff"), +) ANIMATED_MONTAGE_GRID = 3 """Animated images become a grid x grid frame montage (contact sheet).""" @@ -385,36 +392,38 @@ def detect_image_mime_type( *, default_mime_type: str | None = "image/jpeg", ) -> str | None: - """Detect an image MIME type from encoded bytes or a local path. + """Detect an image MIME type by sniffing the file header. + + Only the first bytes of the input are read, so detection cost and memory + stay constant regardless of file size. Args: image_source: Encoded image bytes or a local image path to inspect. default_mime_type: MIME type to return when detection fails. Returns: - The detected MIME type, ``application/octet-stream`` for a recognized - format without a registered MIME, or ``default_mime_type`` when detection - fails. + The detected MIME type, or ``default_mime_type`` when the header does + not match a known image format. """ try: - image_file = ( - io.BytesIO(image_source) - if isinstance(image_source, bytes) - else image_source - ) - with PILImage.open(image_file) as image: - image.verify() - image_format = str(image.format or "").upper() - except Exception as exc: - if not is_recoverable_image_error(exc): - raise + if isinstance(image_source, bytes): + header = image_source[:32] + else: + with open(image_source, "rb") as image_file: + header = image_file.read(32) + except OSError: return default_mime_type - # A decoded format must never inherit an unverified input MIME hint. - return IMAGE_FORMAT_MIME_TYPES.get( - image_format, PILImage.MIME.get(image_format, "application/octet-stream") - ) + if len(header) >= 12: + if header[:4] == b"RIFF" and header[8:12] == b"WEBP": + return "image/webp" + if header[4:8] == b"ftyp" and (b"avif" in header[8:] or b"avis" in header[8:]): + return "image/avif" + for magic, mime_type in _IMAGE_MAGIC_MIME_TYPES: + if header.startswith(magic): + return mime_type + return default_mime_type async def detect_image_mime_type_async( @@ -1235,13 +1244,14 @@ def _convert_image_bytes_sync( Returns: Single-frame JPEG or PNG bytes. An oriented JPEG or PNG within the size - limit is reused unchanged; anything else is re-encoded. + and reuse-byte limits is reused unchanged; anything else is re-encoded. """ with PILImage.open(io.BytesIO(source_bytes)) as image: if ( image.format in {"PNG", "JPEG"} and image.getexif().get(274, 1) == 1 and max(image.size) <= max_size + and len(source_bytes) <= MODEL_IMAGE_REUSE_MAX_BYTES ): return source_bytes cache_key = _image_convert_cache_key( @@ -1389,6 +1399,14 @@ async def prepare_model_image( """ try: async with MediaResolver(image_ref, media_type="image").as_path() as source: + input_size = source.path.stat().st_size + if input_size > MODEL_IMAGE_MAX_INPUT_BYTES: + logger.warning( + "Skipping oversized image input (%d bytes): %s", + input_size, + source.path, + ) + return None image_bytes = await asyncio.to_thread(source.read_bytes) frame_count = await asyncio.to_thread(_inspect_image, image_bytes) is_montage = frame_count > 1 diff --git a/astrbot/dashboard/api/chat.py b/astrbot/dashboard/api/chat.py index 2538a06b3..403c26360 100644 --- a/astrbot/dashboard/api/chat.py +++ b/astrbot/dashboard/api/chat.py @@ -16,6 +16,8 @@ from astrbot.dashboard.schemas import ( ChatThreadMessageRequest, ) from astrbot.dashboard.services.chat_service import ( + MAX_UPLOAD_FILE_SIZE_BYTES, + MAX_UPLOAD_FILE_SIZE_MB, ChatService, ChatServiceError, ) @@ -548,6 +550,11 @@ async def dashboard_post_file( service: ChatService = Depends(get_service), ): try: + content_length = int(request.headers.get("content-length") or 0) + if content_length > MAX_UPLOAD_FILE_SIZE_BYTES: + raise ChatServiceError( + f"File too large (limit {MAX_UPLOAD_FILE_SIZE_MB} MB)" + ) upload = await single_upload(request) if upload is None: raise ChatServiceError("Missing key: file") diff --git a/astrbot/dashboard/services/chat_service.py b/astrbot/dashboard/services/chat_service.py index ca651f645..f23bb236d 100644 --- a/astrbot/dashboard/services/chat_service.py +++ b/astrbot/dashboard/services/chat_service.py @@ -36,6 +36,9 @@ from astrbot.core.utils.media_utils import ( SSE_HEARTBEAT = ": heartbeat\n\n" CHAT_RUN_SUBSCRIBER_QUEUE_SIZE = 256 +# Uploaded chat attachments larger than this are rejected. +MAX_UPLOAD_FILE_SIZE_MB = 512 +MAX_UPLOAD_FILE_SIZE_BYTES = MAX_UPLOAD_FILE_SIZE_MB * 1024 * 1024 WEBCHAT_IMAGE_MIME_TYPES = { ".jpg": "image/jpeg", ".jpeg": "image/jpeg", @@ -596,6 +599,10 @@ class ChatService: return await self.resolve_attachment_file(attachment_id) async def save_uploaded_file(self, file) -> dict: + if (file.content_length or 0) > MAX_UPLOAD_FILE_SIZE_BYTES: + raise ChatServiceError( + f"File too large (limit {MAX_UPLOAD_FILE_SIZE_MB} MB)" + ) filename = sanitize_upload_filename(file.filename) content_type = file.content_type or "application/octet-stream" @@ -614,6 +621,11 @@ class ChatService: raise ChatServiceError("Invalid filename") await file.save(str(file_path)) + if file_path.stat().st_size > MAX_UPLOAD_FILE_SIZE_BYTES: + file_path.unlink(missing_ok=True) + raise ChatServiceError( + f"File too large (limit {MAX_UPLOAD_FILE_SIZE_MB} MB)" + ) if attach_type == "image": detected_mime_type = await detect_image_mime_type_async( file_path, diff --git a/tests/test_chat_route.py b/tests/test_chat_route.py index 37881182d..02f749563 100644 --- a/tests/test_chat_route.py +++ b/tests/test_chat_route.py @@ -508,3 +508,39 @@ async def test_chat_stream_forwards_follow_up_status_by_default( run.task.cancel() await asyncio.gather(run.task, return_exceptions=True) chat_service.webchat_queue_mgr.remove_queues(session_id) + + +@pytest.mark.asyncio +async def test_save_uploaded_file_rejects_oversized_content_length( + chat_service_instance, +): + class FakeUpload: + filename = "big.bin" + content_type = "application/octet-stream" + content_length = chat_service.MAX_UPLOAD_FILE_SIZE_BYTES + 1 + + with pytest.raises(ChatServiceError, match="File too large"): + await chat_service_instance.save_uploaded_file(FakeUpload()) + + +@pytest.mark.asyncio +async def test_save_uploaded_file_rejects_oversized_saved_file( + chat_service_instance, +): + from pathlib import Path + + class FakeUpload: + filename = "big.bin" + content_type = "application/octet-stream" + content_length = None + + async def save(self, path): + saved = Path(path) + saved.write_bytes(b"x") + with saved.open("rb+") as f: + f.truncate(chat_service.MAX_UPLOAD_FILE_SIZE_BYTES + 1) + + with pytest.raises(ChatServiceError, match="File too large"): + await chat_service_instance.save_uploaded_file(FakeUpload()) + + assert not list(Path(chat_service_instance.attachments_dir).iterdir()) diff --git a/tests/test_dashboard.py b/tests/test_dashboard.py index 03dd7a7ce..583023d04 100644 --- a/tests/test_dashboard.py +++ b/tests/test_dashboard.py @@ -10,7 +10,7 @@ import zipfile from datetime import datetime from pathlib import Path from types import SimpleNamespace -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock from urllib.parse import parse_qs, urlsplit, urlunsplit import jwt @@ -2560,7 +2560,7 @@ async def test_plugins( core_lifecycle_td: AstrBotCoreLifecycle, monkeypatch, ): - """测试插件 API 端点，使用 Mock 避免真实网络调用。""" + """Test plugin endpoints with mocked marketplace and repository operations.""" test_client = app.test_client() # 已经安装的插件 @@ -2577,7 +2577,15 @@ async def test_plugins( assert isinstance(installed_at, str) datetime.fromisoformat(installed_at) - # 插件市场 + # Keep marketplace requests independent of external services and local caches. + market_plugins = { + "test/market-plugin": { + "name": "market-plugin", + "repo": "https://github.com/test/market-plugin", + } + } + get_online_plugins = AsyncMock(return_value=(market_plugins, None)) + monkeypatch.setattr(PluginService, "get_online_plugins", get_online_plugins) response = await test_client.get( "/api/plugin/market_list", headers=authenticated_header, @@ -2585,6 +2593,10 @@ async def test_plugins( assert response.status_code == 200 data = await response.get_json() assert data["status"] == "ok" + assert data["data"] == market_plugins + get_online_plugins.assert_awaited_once_with( + custom_registry=None, force_refresh=False + ) # 使用 MockPluginBuilder 创建测试插件 plugin_store_path = core_lifecycle_td.plugin_manager.plugin_store_path diff --git a/tests/test_media_utils.py b/tests/test_media_utils.py index 976efc6db..60371c58c 100644 --- a/tests/test_media_utils.py +++ b/tests/test_media_utils.py @@ -207,6 +207,41 @@ def test_detect_image_mime_type_accepts_path(tmp_path): ) +def test_detect_image_mime_type_sniffs_common_headers(): + """Header sniffing recognizes common formats from bytes alone.""" + assert ( + media_utils.detect_image_mime_type(b"\x89PNG\r\n\x1a\n" + b"\x00" * 24) + == "image/png" + ) + assert ( + media_utils.detect_image_mime_type(b"\xff\xd8\xff\xe0" + b"\x00" * 28) + == "image/jpeg" + ) + assert media_utils.detect_image_mime_type(b"GIF89a" + b"\x00" * 26) == "image/gif" + assert ( + media_utils.detect_image_mime_type(b"RIFF\x00\x00\x00\x00WEBPVP8 ") + == "image/webp" + ) + assert ( + media_utils.detect_image_mime_type(b"\x00\x00\x00\x20ftypavif" + b"\x00" * 20) + == "image/avif" + ) + + +def test_detect_image_mime_type_returns_default_for_unknown_input(): + """Unknown or empty headers fall back to the provided default.""" + assert ( + media_utils.detect_image_mime_type( + b"definitely not an image", default_mime_type=None + ) + is None + ) + assert ( + media_utils.detect_image_mime_type(b"", default_mime_type="image/jpeg") + == "image/jpeg" + ) + + @pytest.mark.asyncio async def test_resolve_image_ref_to_base64_data_decodes_data_uri(tmp_path, monkeypatch): from PIL import Image as PILImage @@ -873,3 +908,53 @@ async def test_wav_to_tencent_silk_skips_resample_for_supported_rate( assert len(fake.calls) == 1 assert fake.calls[0]["sample_rate"] == 24000 + + +@pytest.mark.asyncio +async def test_prepare_model_image_skips_oversized_input(tmp_path, monkeypatch): + """Inputs above the model-image byte cap must be skipped before decoding.""" + from PIL import Image as PILImage + + monkeypatch.setattr(media_utils, "get_astrbot_temp_path", lambda: str(tmp_path)) + image_path = tmp_path / "oversized.png" + PILImage.new("RGB", (4, 4)).save(image_path, format="PNG") + with image_path.open("ab") as f: + f.truncate(media_utils.MODEL_IMAGE_MAX_INPUT_BYTES + 1) + + result = await media_utils.prepare_model_image( + str(image_path), max_size=1280, output_dir=tmp_path + ) + + assert result is None + + +def test_convert_image_bytes_reuses_small_in_range_input(): + """A small oriented in-range PNG keeps its original bytes.""" + from PIL import Image as PILImage + + buffer = BytesIO() + PILImage.new("RGB", (10, 10), (255, 0, 0)).save(buffer, format="PNG") + source = buffer.getvalue() + + result = media_utils._convert_image_bytes_sync(source, 1280, 95) + + assert result is source + + +def test_convert_image_bytes_reencodes_large_in_range_input(tmp_path, monkeypatch): + """An in-range but byte-heavy PNG is re-encoded instead of reused.""" + from PIL import Image as PILImage + + monkeypatch.setattr(media_utils, "get_astrbot_temp_path", lambda: str(tmp_path)) + img = PILImage.new("RGB", (1024, 1024)) + img.frombytes(os.urandom(1024 * 1024 * 3)) + buffer = BytesIO() + img.save(buffer, format="PNG") + source = buffer.getvalue() + assert len(source) > media_utils.MODEL_IMAGE_REUSE_MAX_BYTES + + result = media_utils._convert_image_bytes_sync(source, 1280, 95) + + assert result is not source + assert result[:2] == b"\xff\xd8" + assert len(result) < len(source) diff --git a/tests/test_openai_source.py b/tests/test_openai_source.py index e92d90f8d..b7a9ee5a7 100644 --- a/tests/test_openai_source.py +++ b/tests/test_openai_source.py @@ -5,6 +5,7 @@ from types import SimpleNamespace import httpx import pytest +from openai import AsyncAzureOpenAI, AsyncOpenAI from openai.types.chat.chat_completion import ChatCompletion from openai.types.chat.chat_completion_chunk import ChatCompletionChunk from PIL import Image as PILImage @@ -62,6 +63,60 @@ def _make_groq_provider(overrides: dict | None = None) -> ProviderGroq: ) +@pytest.mark.parametrize( + ("overrides", "expected_client"), + [ + ({}, AsyncOpenAI), + ( + { + "api_version": "2024-02-01", + "api_base": "https://example.openai.azure.com/openai", + }, + AsyncAzureOpenAI, + ), + ], +) +@pytest.mark.asyncio +async def test_provider_client_disables_sdk_builtin_retries(overrides, expected_client): + provider = _make_provider(overrides) + try: + assert isinstance(provider.client, expected_client) + assert provider.client.max_retries == 0 + finally: + await provider.terminate() + + +@pytest.mark.asyncio +async def test_query_attempts_exactly_request_max_retries_times(monkeypatch): + monkeypatch.setattr(request_retry, "REQUEST_RETRY_WAIT_MIN_S", 0) + monkeypatch.setattr(request_retry, "REQUEST_RETRY_WAIT_MAX_S", 0) + + provider = _make_provider() + try: + calls = 0 + + async def failing_create(**kwargs): + nonlocal calls + calls += 1 + raise httpx.ConnectError("temporary connection failure") + + monkeypatch.setattr(provider.client.chat.completions, "create", failing_create) + + with pytest.raises(httpx.ConnectError): + await provider._query( + payloads={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}], + }, + tools=None, + request_max_retries=2, + ) + + assert calls == 2 + finally: + await provider.terminate() + + def test_create_http_client_uses_openai_httpx_module(monkeypatch): captured: dict[str, object] = {} fake_httpx_module = object() diff --git a/tests/test_process_stage_images.py b/tests/test_process_stage_images.py index eaa39c4bc..86f109a02 100644 --- a/tests/test_process_stage_images.py +++ b/tests/test_process_stage_images.py @@ -4,7 +4,6 @@ import asyncio import base64 import copy import inspect -import os from contextlib import nullcontext from pathlib import Path from types import SimpleNamespace @@ -271,44 +270,6 @@ async def test_cua_runtime_keeps_input_image_geometry( assert Path(req.image_urls[0]).read_bytes() == original -@pytest.mark.asyncio -@pytest.mark.parametrize( - ("runtime", "booter", "big"), - [ - ("sandbox", "cua", True), - ("sandbox", "cua", False), - ("sandbox", "shipyard_neo", True), - ], -) -async def test_cua_oversize_image_warns( - harness, tmp_path, monkeypatch, runtime, booter, big -): - """Byte-exact CUA passthrough warns when an image may exceed upload limits.""" - dims = (1500, 1500) if big else (60, 30) - path = tmp_path / "shot.png" - if big: - # Noise stays incompressible, keeping the PNG above the 5 MB threshold. - PILImage.frombytes("RGB", dims, os.urandom(dims[0] * dims[1] * 3)).save( - path, "PNG" - ) - else: - PILImage.new("RGB", dims, "red").save(path, "PNG") - original = path.read_bytes() - harness.config["provider_settings"]["computer_use_runtime"] = runtime - harness.config["provider_settings"]["sandbox"] = {"booter": booter} - warning = MagicMock() - monkeypatch.setattr(internal.logger, "warning", warning) - event = make_event([Image(file=str(path))], text="") - await process_event(harness, event, preprocess_first=True) - assert len(harness.captured) == 1 - req = harness.captured[0].req - gated = runtime == "sandbox" and booter == "cua" - if gated: - assert Path(req.image_urls[0]).read_bytes() == original - warned = any("upload limits" in str(call) for call in warning.call_args_list) - assert warned == (big and gated) - - @pytest.mark.asyncio async def test_cua_montage_keeps_configured_cap(harness, tmp_path): """Animated inputs keep the configured montage cap under CUA pixel mode.""" diff --git a/tests/unit/test_webchat_upload_image_format.py b/tests/unit/test_webchat_upload_image_format.py index 7125c6961..4938a2f93 100644 --- a/tests/unit/test_webchat_upload_image_format.py +++ b/tests/unit/test_webchat_upload_image_format.py @@ -18,6 +18,7 @@ async def test_webchat_upload_uses_detected_image_type(tmp_path): class FakeUploadFile: filename = "pasted.png" content_type = "image/png" + content_length = None async def save(self, destination): with open(destination, "wb") as output: