From 5dd8142cc0f9c5f2f7b615c0f1cc15a12c047b9c Mon Sep 17 00:00:00 2001 From: SweetenedSuzuka <188338189+SweetenedSuzuka@users.noreply.github.com> Date: Fri, 14 Aug 2026 00:55:39 +0800 Subject: [PATCH 1/7] fix(provider): disable OpenAI SDK built-in retries under AstrBot retry layer (#9663) AsyncOpenAI/AsyncAzureOpenAI inherit the SDK default max_retries=2, which nests underneath AstrBot's tenacity-based retry_provider_request(). As a result provider_settings.request_max_retries no longer reflects the true attempt count (e.g. request_max_retries=1 still sends up to 3 HTTP requests) and delays fallback-provider switchover. Pass max_retries=0 so AstrBot's retry layer is the single source of truth for retries. --- .../core/provider/sources/openai_source.py | 6 ++++++ tests/test_openai_source.py | 21 +++++++++++++++++++ 2 files changed, 27 insertions(+) diff --git a/astrbot/core/provider/sources/openai_source.py b/astrbot/core/provider/sources/openai_source.py index f7870b7137..e984b076b6 100644 --- a/astrbot/core/provider/sources/openai_source.py +++ b/astrbot/core/provider/sources/openai_source.py @@ -377,6 +377,9 @@ def __init__(self, provider_config, provider_settings) -> None: default_headers=self.custom_headers, base_url=provider_config.get("api_base", ""), 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: @@ -386,6 +389,9 @@ def __init__(self, provider_config, provider_settings) -> None: base_url=provider_config.get("api_base", 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/tests/test_openai_source.py b/tests/test_openai_source.py index 911b76131f..b719178e7f 100644 --- a/tests/test_openai_source.py +++ b/tests/test_openai_source.py @@ -5,6 +5,7 @@ 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 @@ -60,6 +61,26 @@ def _make_groq_provider(overrides: dict | None = None) -> ProviderGroq: ) +@pytest.mark.asyncio +async def test_openai_client_disables_sdk_builtin_retries(): + provider = _make_provider() + try: + assert isinstance(provider.client, AsyncOpenAI) + assert provider.client.max_retries == 0 + finally: + await provider.terminate() + + +@pytest.mark.asyncio +async def test_azure_client_disables_sdk_builtin_retries(): + provider = _make_provider({"api_version": "2024-02-01"}) + try: + assert isinstance(provider.client, AsyncAzureOpenAI) + assert provider.client.max_retries == 0 + finally: + await provider.terminate() + + def test_create_http_client_uses_openai_httpx_module(monkeypatch): captured: dict[str, object] = {} fake_httpx_module = object() From 5735f0c281d82b7d5952ad8868f95514458d4c71 Mon Sep 17 00:00:00 2001 From: SweetenedSuzuka <188338189+SweetenedSuzuka@users.noreply.github.com> Date: Fri, 14 Aug 2026 01:10:57 +0800 Subject: [PATCH 2/7] test(provider): add retry-count regression test, parameterize client retry tests Address review feedback on #9669: merge the near-identical official/Azure client construction tests into a parametrized test, and add a regression test asserting _query() performs exactly request_max_retries attempts through retry_provider_request() when the call keeps failing, guarding against multiplicative retries being reintroduced. --- tests/test_openai_source.py | 42 ++++++++++++++++++++++++++++++------- 1 file changed, 35 insertions(+), 7 deletions(-) diff --git a/tests/test_openai_source.py b/tests/test_openai_source.py index b719178e7f..ac36299f23 100644 --- a/tests/test_openai_source.py +++ b/tests/test_openai_source.py @@ -61,22 +61,50 @@ def _make_groq_provider(overrides: dict | None = None) -> ProviderGroq: ) +@pytest.mark.parametrize( + ("overrides", "expected_client"), + [ + ({}, AsyncOpenAI), + ({"api_version": "2024-02-01"}, AsyncAzureOpenAI), + ], +) @pytest.mark.asyncio -async def test_openai_client_disables_sdk_builtin_retries(): - provider = _make_provider() +async def test_provider_client_disables_sdk_builtin_retries(overrides, expected_client): + provider = _make_provider(overrides) try: - assert isinstance(provider.client, AsyncOpenAI) + assert isinstance(provider.client, expected_client) assert provider.client.max_retries == 0 finally: await provider.terminate() @pytest.mark.asyncio -async def test_azure_client_disables_sdk_builtin_retries(): - provider = _make_provider({"api_version": "2024-02-01"}) +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: - assert isinstance(provider.client, AsyncAzureOpenAI) - assert provider.client.max_retries == 0 + 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() From 7ab86dda9b9d513259ce3c0edc4f62d6930819da Mon Sep 17 00:00:00 2001 From: RuochenPan Date: Thu, 17 Sep 2026 17:51:17 +0800 Subject: [PATCH 3/7] fix(media): guard oversized image inputs and attachment uploads MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 为模型图片输入与聊天附件上传增加大小护栏 / Add size guards for model image inputs and chat attachment uploads: - _convert_image_bytes_sync 复用原始编码字节的分支增加 2MB 上限;更大的 输入重编码,输出始终受像素尺寸与质量约束 - prepare_model_image 在读入前检查文件大小,超过 32MB 跳过并记录日志 - 聊天附件上传上限 100MB:路由先检查 Content-Length,服务层落盘后复核 实际大小 - Reuse the original encoded bytes in _convert_image_bytes_sync only for stills up to 2MB; larger inputs are re-encoded so the output stays bounded by pixel size and quality. - prepare_model_image checks the file size before reading and skips inputs above 32MB with a warning log. - Chat attachment uploads are capped at 100MB: the route checks Content-Length first, and save_uploaded_file re-checks the on-disk size after saving. Tests: tests/test_media_utils.py, tests/test_chat_route.py --- astrbot/core/utils/media_utils.py | 16 ++++++- astrbot/dashboard/api/chat.py | 7 +++ astrbot/dashboard/services/chat_service.py | 12 ++++++ tests/test_chat_route.py | 36 ++++++++++++++++ tests/test_media_utils.py | 50 ++++++++++++++++++++++ 5 files changed, 120 insertions(+), 1 deletion(-) diff --git a/astrbot/core/utils/media_utils.py b/astrbot/core/utils/media_utils.py index ff552293d5..b527c4b6c2 100644 --- a/astrbot/core/utils/media_utils.py +++ b/astrbot/core/utils/media_utils.py @@ -42,6 +42,11 @@ 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", @@ -1235,13 +1240,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( @@ -1388,6 +1394,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) if frame_count > 1: diff --git a/astrbot/dashboard/api/chat.py b/astrbot/dashboard/api/chat.py index 2538a06b3b..403c263602 100644 --- a/astrbot/dashboard/api/chat.py +++ b/astrbot/dashboard/api/chat.py @@ -16,6 +16,8 @@ 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 ca651f6452..b94443b1be 100644 --- a/astrbot/dashboard/services/chat_service.py +++ b/astrbot/dashboard/services/chat_service.py @@ -36,6 +36,9 @@ SSE_HEARTBEAT = ": heartbeat\n\n" CHAT_RUN_SUBSCRIBER_QUEUE_SIZE = 256 +# Uploaded chat attachments larger than this are rejected. +MAX_UPLOAD_FILE_SIZE_MB = 100 +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 @@ async def resolve_attachment_file_from_dashboard_query( 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 @@ async def save_uploaded_file(self, file) -> dict: 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 37881182db..02f749563b 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_media_utils.py b/tests/test_media_utils.py index 976efc6db3..46dd1f66b4 100644 --- a/tests/test_media_utils.py +++ b/tests/test_media_utils.py @@ -873,3 +873,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) From f8eb062c880f4e2c1a105642a75c791b831e8982 Mon Sep 17 00:00:00 2001 From: zenfun Date: Thu, 17 Sep 2026 16:59:53 +0800 Subject: [PATCH 4/7] perf: reduce image memory overhead across lifecycle --- .github/workflows/image-memory-validation.yml | 215 ++++ astrbot/core/agent/context/compressor.py | 31 + astrbot/core/agent/context/config.py | 3 + astrbot/core/agent/context/image_budget.py | 87 ++ astrbot/core/agent/context/manager.py | 5 + astrbot/core/agent/context/token_counter.py | 11 +- astrbot/core/agent/message.py | 22 + .../agent/runners/coze/coze_agent_runner.py | 69 +- .../runners/deerflow/deerflow_agent_runner.py | 21 +- .../deerflow/deerflow_content_mapper.py | 12 +- .../agent/runners/dify/dify_agent_runner.py | 22 +- .../agent/runners/tool_loop_agent_runner.py | 128 +- astrbot/core/astr_main_agent.py | 27 +- astrbot/core/computer/file_read_utils.py | 82 +- astrbot/core/config/default.py | 10 + astrbot/core/exceptions.py | 4 + .../method/agent_sub_stages/internal.py | 6 +- .../method/agent_sub_stages/third_party.py | 16 +- astrbot/core/provider/entities.py | 42 +- astrbot/core/provider/modalities.py | 8 +- .../core/provider/sources/anthropic_source.py | 43 +- .../core/provider/sources/gemini_source.py | 56 +- .../sources/openai_responses_source.py | 22 + .../core/provider/sources/openai_source.py | 82 +- .../core/provider/sources/request_retry.py | 14 +- astrbot/core/tools/computer_tools/fs.py | 2 + astrbot/core/utils/image_media_store.py | 357 ++++++ astrbot/core/utils/media_utils.py | 1109 +++++++++++++++-- astrbot/dashboard/api/conversations.py | 17 +- .../services/conversation_service.py | 85 +- .../src/api/generated/openapi-v1/sdk.gen.ts | 12 +- .../src/api/generated/openapi-v1/types.gen.ts | 14 + .../chat/AuthenticatedMediaImage.vue | 70 ++ .../ConversationHistoryPreview.vue | 39 +- .../locales/en-US/features/conversation.json | 2 + .../locales/ja-JP/features/conversation.json | 2 + .../locales/ru-RU/features/conversation.json | 2 + .../locales/zh-CN/features/conversation.json | 2 + .../ConversationWorkspacePage.vue | 2 + docs/en/dev/openapi-scopes.md | 1 + docs/public/openapi.json | 49 + docs/zh/dev/openapi-scopes.md | 1 + .../.openspec.yaml | 2 + .../optimize-image-memory-lifecycle/design.md | 152 +++ .../proposal.md | 28 + .../specs/agent-context-image-budget/spec.md | 40 + .../specs/conversation-history-media/spec.md | 37 + .../specs/image-memory-lifecycle/spec.md | 26 + .../optimize-image-memory-lifecycle/tasks.md | 36 + openspec/openapi-v1.yaml | 30 + scripts/image_memory_bench/ablation_runner.py | 66 + .../image_memory_bench/baseline_manifest.py | 54 + .../image_memory_bench/generate_fixtures.py | 127 ++ .../image_memory_bench/history_benchmark.py | 461 +++++++ .../image_memory_bench/lifecycle_workloads.py | 371 ++++++ scripts/image_memory_bench/measure.py | 399 ++++++ scripts/image_memory_bench/migrate_history.py | 99 ++ .../prepare_lifecycle_fixture.py | 131 ++ scripts/manage_image_history.py | 434 +++++++ tests/test_agent_runner_media_resolver.py | 114 ++ tests/test_computer_fs_tools.py | 49 +- tests/test_fastapi_v1_dashboard.py | 40 + tests/test_media_utils.py | 26 +- tests/test_openai_responses_source.py | 29 + tests/test_openai_source.py | 42 +- tests/test_process_stage_images.py | 21 +- tests/test_tool_loop_agent_runner.py | 341 ++++- tests/unit/test_astr_main_agent.py | 23 + tests/unit/test_context_image_budget.py | 153 +++ tests/unit/test_conversation_media_api.py | 149 +++ tests/unit/test_image_lifecycle_workloads.py | 36 + tests/unit/test_image_media_store.py | 304 +++++ tests/unit/test_image_memory_benchmark.py | 59 + tests/unit/test_image_preparation_budget.py | 376 ++++++ tests/unit/test_image_screenshot_encoding.py | 43 + tests/unit/test_image_source_integration.py | 116 ++ tests/unit/test_image_source_preparation.py | 134 ++ tests/unit/test_manage_image_history.py | 354 ++++++ tests/unit/test_provider_image_references.py | 191 +++ tests/unit/test_provider_request_retry.py | 40 + 80 files changed, 7664 insertions(+), 273 deletions(-) create mode 100644 .github/workflows/image-memory-validation.yml create mode 100644 astrbot/core/agent/context/image_budget.py create mode 100644 astrbot/core/utils/image_media_store.py create mode 100644 dashboard/src/components/chat/AuthenticatedMediaImage.vue create mode 100644 openspec/changes/optimize-image-memory-lifecycle/.openspec.yaml create mode 100644 openspec/changes/optimize-image-memory-lifecycle/design.md create mode 100644 openspec/changes/optimize-image-memory-lifecycle/proposal.md create mode 100644 openspec/changes/optimize-image-memory-lifecycle/specs/agent-context-image-budget/spec.md create mode 100644 openspec/changes/optimize-image-memory-lifecycle/specs/conversation-history-media/spec.md create mode 100644 openspec/changes/optimize-image-memory-lifecycle/specs/image-memory-lifecycle/spec.md create mode 100644 openspec/changes/optimize-image-memory-lifecycle/tasks.md create mode 100644 scripts/image_memory_bench/ablation_runner.py create mode 100644 scripts/image_memory_bench/baseline_manifest.py create mode 100644 scripts/image_memory_bench/generate_fixtures.py create mode 100644 scripts/image_memory_bench/history_benchmark.py create mode 100644 scripts/image_memory_bench/lifecycle_workloads.py create mode 100644 scripts/image_memory_bench/measure.py create mode 100644 scripts/image_memory_bench/migrate_history.py create mode 100644 scripts/image_memory_bench/prepare_lifecycle_fixture.py create mode 100644 scripts/manage_image_history.py create mode 100644 tests/unit/test_context_image_budget.py create mode 100644 tests/unit/test_conversation_media_api.py create mode 100644 tests/unit/test_image_lifecycle_workloads.py create mode 100644 tests/unit/test_image_media_store.py create mode 100644 tests/unit/test_image_memory_benchmark.py create mode 100644 tests/unit/test_image_preparation_budget.py create mode 100644 tests/unit/test_image_screenshot_encoding.py create mode 100644 tests/unit/test_image_source_integration.py create mode 100644 tests/unit/test_image_source_preparation.py create mode 100644 tests/unit/test_manage_image_history.py create mode 100644 tests/unit/test_provider_image_references.py create mode 100644 tests/unit/test_provider_request_retry.py diff --git a/.github/workflows/image-memory-validation.yml b/.github/workflows/image-memory-validation.yml new file mode 100644 index 0000000000..a1ff61278c --- /dev/null +++ b/.github/workflows/image-memory-validation.yml @@ -0,0 +1,215 @@ +name: Image Memory Lifecycle Validation + +on: + push: + branches: + - feat/optimize-image-memory-lifecycle + workflow_dispatch: + +jobs: + validate: + name: Image and memory validation (${{ matrix.os }}) + runs-on: ${{ matrix.os }} + timeout-minutes: 90 + strategy: + fail-fast: false + matrix: + os: [ubuntu-latest, windows-latest, macos-latest] + + defaults: + run: + shell: bash + + steps: + - name: Checkout + uses: actions/checkout@v7 + + - name: Set up Python + uses: actions/setup-python@v7 + with: + python-version: '3.12' + + - name: Install uv + run: | + python -m pip install --upgrade pip + python -m pip install uv + + - name: Install dependencies + run: uv sync + + - name: Run repository ruff checks + run: | + uv run ruff format --check . + uv run ruff check . + + - name: Run focused tests + run: | + mkdir -p "$RUNNER_TEMP/image-memory-validation/results" + uv run pytest -q \ + tests/unit/test_image_preparation_budget.py \ + tests/unit/test_image_media_store.py \ + tests/test_media_utils.py \ + tests/test_computer_fs_tools.py \ + tests/test_tool_loop_agent_runner.py \ + tests/agent/test_context_manager.py \ + tests/unit/test_astr_main_agent.py \ + tests/unit/test_context_image_budget.py \ + tests/unit/test_image_source_preparation.py \ + tests/unit/test_image_memory_benchmark.py \ + tests/unit/test_image_lifecycle_workloads.py \ + tests/unit/test_image_screenshot_encoding.py \ + tests/unit/test_image_source_integration.py \ + tests/unit/test_provider_image_references.py \ + tests/test_agent_runner_media_resolver.py \ + tests/unit/test_provider_request_retry.py \ + tests/unit/test_manage_image_history.py \ + tests/unit/test_conversation_media_api.py \ + tests/test_openai_source.py \ + tests/test_openai_responses_source.py \ + tests/test_gemini_source.py \ + tests/test_anthropic_source.py \ + tests/test_anthropic_kimi_code_provider.py \ + tests/test_deerflow_agent_runner.py \ + tests/unit/test_cron_context_compression.py \ + tests/unit/test_group_chat_context_wiring.py \ + tests/test_platform_image_format_preservation.py \ + tests/test_webchat_queue_lifecycle.py \ + tests/test_webchat_upload_image_format.py \ + --junitxml="$RUNNER_TEMP/image-memory-validation/results/tests.xml" + + - name: Generate reproducible fixtures outside measured processes + run: | + mkdir -p "$RUNNER_TEMP/image-memory-validation/fixtures" + uv run python scripts/image_memory_bench/generate_fixtures.py \ + "$RUNNER_TEMP/image-memory-validation/fixtures" \ + --seeds 7 19 43 + + - name: Record versions and fixture manifest + run: | + mkdir -p "$RUNNER_TEMP/image-memory-validation/results" + uv run python scripts/image_memory_bench/baseline_manifest.py \ + "$RUNNER_TEMP/image-memory-validation/results/manifest.json" \ + --fixtures "$RUNNER_TEMP/image-memory-validation/fixtures" + + - name: Run single-image memory workloads for every paired seed + run: | + set +e + result="$RUNNER_TEMP/image-memory-validation/results/measurements.jsonl" + failures=0 + for seed in 7 19 43; do + while IFS= read -r image; do + uv run python scripts/image_memory_bench/measure.py \ + "$image" "$result" \ + --repo "$GITHUB_WORKSPACE" \ + --csv "$RUNNER_TEMP/image-memory-validation/results/measurements.csv" \ + --workload "$seed-$(basename "$image")" + code=$? + if [ "$code" -ne 0 ]; then + failures=$((failures + 1)) + fi + done < <(find "$RUNNER_TEMP/image-memory-validation/fixtures/$seed" \ + -type f \( -name '*ordinary*' -o -name '*stress*' \) | sort) + done + printf '{"measurement_failures":%s}\n' "$failures" \ + > "$RUNNER_TEMP/image-memory-validation/results/measurement-summary.json" + + - name: Run Python allocation diagnostic workload + run: | + uv run python scripts/image_memory_bench/measure.py \ + "$RUNNER_TEMP/image-memory-validation/fixtures/7/png-stress.png" \ + "$RUNNER_TEMP/image-memory-validation/results/trace-python.jsonl" \ + --repo "$GITHUB_WORKSPACE" \ + --trace-python \ + --csv "$RUNNER_TEMP/image-memory-validation/results/trace-python.csv" \ + --workload trace-python-png-stress + + - name: Run paired history loading workloads + run: | + mkdir -p "$RUNNER_TEMP/image-memory-validation/results/history" + for seed in 7 19 43; do + uv run python scripts/image_memory_bench/history_benchmark.py \ + "$RUNNER_TEMP/image-memory-validation/fixtures/$seed/manifest.json" \ + "$RUNNER_TEMP/image-memory-validation/results/history/$seed.jsonl" \ + --repo "$GITHUB_WORKSPACE" \ + --images 50 \ + --keep-turns 25 \ + --repeats 10 \ + --cold-repeats 10 \ + --warm-repeats 10 \ + --distinct \ + --csv "$RUNNER_TEMP/image-memory-validation/results/history/$seed.csv" \ + --rss-jsonl "$RUNNER_TEMP/image-memory-validation/results/history/$seed.rss.jsonl" + done + + - name: Run fixed-window concurrent lifecycle and restart workload + run: | + lifecycle_root="$RUNNER_TEMP/image-memory-validation/lifecycle" + mkdir -p "$lifecycle_root" + uv run python scripts/image_memory_bench/prepare_lifecycle_fixture.py \ + "$RUNNER_TEMP/image-memory-validation/fixtures/7/manifest.json" \ + "$lifecycle_root/input" \ + --session-count 4 \ + --history-turns 50 + lifecycle_manifest="$lifecycle_root/input/lifecycle-manifest.json" + database=$(uv run python -c \ + 'import json,sys; print(json.load(open(sys.argv[1]))["database"])' \ + "$lifecycle_manifest") + media_root=$(uv run python -c \ + 'import json,sys; print(json.load(open(sys.argv[1]))["media_root"])' \ + "$lifecycle_manifest") + conversation_ids=$(uv run python -c \ + 'import json,sys; print(" ".join(json.load(open(sys.argv[1]))["conversation_ids"]))' \ + "$lifecycle_manifest") + command=( + uv run python scripts/image_memory_bench/lifecycle_workloads.py + "$lifecycle_root/result.json" + --repo "$GITHUB_WORKSPACE" + --database "$database" + --media-root "$media_root" + --session-count 4 + --request-count 200 + --window-turns 25 + --rss-jsonl "$lifecycle_root/result.rss.jsonl" + ) + for conversation_id in $conversation_ids; do + command+=(--conversation-id "$conversation_id") + done + "${command[@]}" + + - name: Record ablation contract status without fabricating factors + run: | + uv run python scripts/image_memory_bench/ablation_runner.py \ + "$RUNNER_TEMP/image-memory-validation/results/ablation-plan.json" \ + --baseline-repo "$GITHUB_WORKSPACE" \ + --candidate-repo "$GITHUB_WORKSPACE" + + - name: Upload raw validation artifacts + if: always() + uses: actions/upload-artifact@v4 + with: + name: image-memory-validation-${{ matrix.os }} + path: ${{ runner.temp }}/image-memory-validation + if-no-files-found: error + + dashboard: + name: Dashboard typecheck + runs-on: ubuntu-latest + timeout-minutes: 20 + defaults: + run: + working-directory: dashboard + steps: + - name: Checkout + uses: actions/checkout@v7 + - name: Set up Node + uses: actions/setup-node@v7 + with: + node-version: 22 + cache: pnpm + cache-dependency-path: dashboard/pnpm-lock.yaml + - name: Enable pnpm + run: corepack enable + - name: Install dashboard dependencies + run: pnpm install --frozen-lockfile + - name: Run dashboard typecheck + run: pnpm typecheck diff --git a/astrbot/core/agent/context/compressor.py b/astrbot/core/agent/context/compressor.py index 759604dd93..7aaa7adf50 100644 --- a/astrbot/core/agent/context/compressor.py +++ b/astrbot/core/agent/context/compressor.py @@ -1,10 +1,20 @@ +from pathlib import Path from typing import TYPE_CHECKING, Protocol, runtime_checkable +from astrbot.core.exceptions import ProviderRequestTooLargeError +from astrbot.core.utils.astrbot_path import get_astrbot_data_path +from astrbot.core.utils.image_media_store import ( + ImageMediaStore, + materialize_image_media_refs, +) +from astrbot.core.utils.media_utils import ImagePayloadTooLargeError + from ...provider.modalities import ( log_context_sanitize_stats, sanitize_contexts_by_modalities, ) from ..message import Message +from .image_budget import get_image_encoded_byte_limit, validate_context_image_bytes from .token_counter import EstimateTokenCounter, TokenCounter if TYPE_CHECKING: @@ -130,6 +140,7 @@ def __init__( instruction_text: str | None = None, compression_threshold: float = 0.82, token_counter: TokenCounter | None = None, + image_media_store: ImageMediaStore | None = None, ) -> None: """Initialize the LLM summary compressor. @@ -139,8 +150,10 @@ def __init__( exact context. Clamped to 0-0.3. instruction_text: Custom instruction for summary generation. compression_threshold: The compression trigger threshold (default: 0.82). + image_media_store: Store shared with the main request. """ self.provider = provider + self.image_media_store = image_media_store self.keep_recent_ratio = min(max(float(keep_recent_ratio), 0.0), 0.3) self.compression_threshold = compression_threshold self.token_counter = token_counter or EstimateTokenCounter() @@ -273,11 +286,29 @@ async def __call__(self, messages: list[Message]) -> list[Message]: # Generate summary try: + limit = get_image_encoded_byte_limit( + getattr(self.provider, "provider_settings", {}) + ) + validate_context_image_bytes(sanitized_summary_contexts, limit) + sanitized_summary_contexts = await materialize_image_media_refs( + sanitized_summary_contexts, + self.image_media_store + or ImageMediaStore(Path(get_astrbot_data_path()) / "media"), + ) response = await self.provider.text_chat( contexts=sanitized_summary_contexts, ) summary_content = (response.completion_text or "").strip() + except (MemoryError, ImagePayloadTooLargeError, ProviderRequestTooLargeError): + raise except Exception as e: + status_code = getattr(e, "status_code", None) or getattr( + getattr(e, "response", None), "status_code", None + ) + if status_code == 413: + raise ProviderRequestTooLargeError( + "The summary provider rejected the request size (HTTP 413)." + ) from e logger.error(f"Failed to generate summary: {e}") return messages diff --git a/astrbot/core/agent/context/config.py b/astrbot/core/agent/context/config.py index aa216d9a25..e2679912b7 100644 --- a/astrbot/core/agent/context/config.py +++ b/astrbot/core/agent/context/config.py @@ -6,6 +6,7 @@ if TYPE_CHECKING: from astrbot.core.provider.provider import Provider + from astrbot.core.utils.image_media_store import ImageMediaStore @dataclass @@ -33,3 +34,5 @@ class ContextConfig: """Custom token counting method. If None, the default method is used.""" custom_compressor: ContextCompressor | None = None """Custom context compression method. If None, the default method is used.""" + image_media_store: "ImageMediaStore | None" = None + """Application media store shared by main and summary requests.""" diff --git a/astrbot/core/agent/context/image_budget.py b/astrbot/core/agent/context/image_budget.py new file mode 100644 index 0000000000..d02eb807ff --- /dev/null +++ b/astrbot/core/agent/context/image_budget.py @@ -0,0 +1,87 @@ +"""Validate image bytes independently from text-token estimates.""" + +from collections.abc import Sequence + +from astrbot.core.agent.message import ImageMediaRefPart, ImageURLPart, Message +from astrbot.core.utils.media_utils import ( + IMAGE_COMPRESS_DEFAULT_MAX_ENCODED_BYTES, + ImagePayloadTooLargeError, +) + + +def get_image_encoded_byte_limit(provider_settings: dict | None) -> int: + """Read the configured per-image limit with the same rules for every request. + + Args: + provider_settings: Settings belonging to the actual request provider. + + Returns: + A positive byte limit, or the default when the setting is invalid. + """ + options = ( + provider_settings.get("image_compress_options", {}) + if isinstance(provider_settings, dict) + else {} + ) + limit = options.get("max_encoded_bytes") if isinstance(options, dict) else None + if isinstance(limit, int) and not isinstance(limit, bool) and limit > 0: + return limit + return IMAGE_COMPRESS_DEFAULT_MAX_ENCODED_BYTES + + +def validate_context_image_bytes( + messages: Sequence[Message | dict], + max_encoded_bytes: int = IMAGE_COMPRESS_DEFAULT_MAX_ENCODED_BYTES, +) -> int: + """Validate selected images without loading or rewriting historical bytes. + + Args: + messages: Selected main or summary request messages. + max_encoded_bytes: Maximum Base64 bytes for a single image. + + Returns: + Total known encoded-image bytes, excluding JSON and data URI headers. + Remote URLs have unknown size until resolved and are not counted. + + Raises: + ImagePayloadTooLargeError: A selected image exceeds the single-image cap. + ValueError: The configured cap is invalid. + """ + if isinstance(max_encoded_bytes, bool) or max_encoded_bytes < 1: + raise ValueError("Image byte budget must be a positive integer") + total = 0 + for message in messages: + parts = ( + message.content if isinstance(message, Message) else message.get("content") + ) + if not isinstance(parts, list): + continue + for part in parts: + size = 0 + url = None + if isinstance(part, ImageMediaRefPart): + size = 4 * ((part.byte_size + 2) // 3) + elif isinstance(part, ImageURLPart): + url = part.image_url.url + elif isinstance(part, dict): + if part.get("type") == "image_media_ref": + size = 4 * ((int(part["byte_size"]) + 2) // 3) + elif part.get("type") == "image_url": + image_url = part.get("image_url") + url = ( + image_url.get("url") + if isinstance(image_url, dict) + else image_url + ) + if isinstance(url, str) and url.startswith("data:image/"): + comma = url.find(",") + if comma >= 0 and ";base64" in url[:comma]: + # Do not slice the full Base64 suffix merely to count it. + size = len(url) - comma - 1 + if size > max_encoded_bytes: + raise ImagePayloadTooLargeError( + f"A selected image uses {size} Base64 bytes, exceeding the " + f"{max_encoded_bytes}-byte limit. Historical images were not changed." + ) + total += size + return total diff --git a/astrbot/core/agent/context/manager.py b/astrbot/core/agent/context/manager.py index 1a11ebff96..fe98f3b7b2 100644 --- a/astrbot/core/agent/context/manager.py +++ b/astrbot/core/agent/context/manager.py @@ -1,4 +1,6 @@ from astrbot import logger +from astrbot.core.exceptions import ProviderRequestTooLargeError +from astrbot.core.utils.media_utils import ImagePayloadTooLargeError from ..message import Message from .compressor import LLMSummaryCompressor, TruncateByTurnsCompressor @@ -36,6 +38,7 @@ def __init__( keep_recent_ratio=config.llm_compress_keep_recent_ratio, instruction_text=config.llm_compress_instruction, token_counter=self.token_counter, + image_media_store=config.image_media_store, ) else: self.compressor = TruncateByTurnsCompressor( @@ -76,6 +79,8 @@ async def process( result = await self._run_compression(result, total_tokens) return result + except (MemoryError, ImagePayloadTooLargeError, ProviderRequestTooLargeError): + raise except Exception as e: logger.error(f"Error during context processing: {e}", exc_info=True) return messages diff --git a/astrbot/core/agent/context/token_counter.py b/astrbot/core/agent/context/token_counter.py index 7c60cb23ec..c93e958c87 100644 --- a/astrbot/core/agent/context/token_counter.py +++ b/astrbot/core/agent/context/token_counter.py @@ -1,7 +1,14 @@ import json from typing import Protocol, runtime_checkable -from ..message import AudioURLPart, ImageURLPart, Message, TextPart, ThinkPart +from ..message import ( + AudioURLPart, + ImageMediaRefPart, + ImageURLPart, + Message, + TextPart, + ThinkPart, +) @runtime_checkable @@ -60,7 +67,7 @@ def count_tokens( total += self._estimate_tokens(part.text) elif isinstance(part, ThinkPart): total += self._estimate_tokens(part.think) - elif isinstance(part, ImageURLPart): + elif isinstance(part, (ImageURLPart, ImageMediaRefPart)): total += IMAGE_TOKEN_ESTIMATE elif isinstance(part, AudioURLPart): total += AUDIO_TOKEN_ESTIMATE diff --git a/astrbot/core/agent/message.py b/astrbot/core/agent/message.py index 4292f4c04e..a754692786 100644 --- a/astrbot/core/agent/message.py +++ b/astrbot/core/agent/message.py @@ -120,11 +120,33 @@ class ImageURL(BaseModel): """The URL of the image, can be data URI scheme like `data:image/png;base64,...`.""" id: str | None = None """The ID of the image, to allow LLMs to distinguish different images.""" + detail: str | None = None + + @model_serializer(mode="wrap") + def serialize(self, handler): + data = handler(self) + if self.detail is None: + data.pop("detail", None) + return data type: str = "image_url" image_url: ImageURL +class ImageMediaRefPart(ContentPart): + """A durable image reference that is resolved only for a provider request.""" + + type: str = "image_media_ref" + media_id: str + mime_type: str + width: int | None = None + height: int | None = None + byte_size: int + detail: str | None = None + version: int = 1 + image_id: str | None = None + + class AudioURLPart(ContentPart): """ >>> AudioURLPart(audio_url=AudioURLPart.AudioURL(url="https://example.com/audio.mp3")).model_dump() diff --git a/astrbot/core/agent/runners/coze/coze_agent_runner.py b/astrbot/core/agent/runners/coze/coze_agent_runner.py index 801031d963..c505d953d6 100644 --- a/astrbot/core/agent/runners/coze/coze_agent_runner.py +++ b/astrbot/core/agent/runners/coze/coze_agent_runner.py @@ -1,16 +1,29 @@ import json import sys import typing as T +from dataclasses import replace +from pathlib import Path import astrbot.core.message.components as Comp from astrbot import logger from astrbot.core import sp +from astrbot.core.agent.context.image_budget import validate_context_image_bytes from astrbot.core.message.message_event_result import MessageChain from astrbot.core.provider.entities import ( LLMResponse, ProviderRequest, ) -from astrbot.core.utils.media_utils import MediaResolver, describe_media_ref +from astrbot.core.utils.astrbot_path import get_astrbot_data_path +from astrbot.core.utils.image_media_store import ( + ImageMediaStore, + materialize_image_media_refs, +) +from astrbot.core.utils.media_utils import ( + ImagePayloadTooLargeError, + ImagePreparationOptions, + describe_media_ref, + resolve_media_ref_to_base64_data, +) from ...hooks import BaseAgentRunHooks from ...message import is_checkpoint_message @@ -91,6 +104,11 @@ async def step(self): # 执行 Coze 请求并处理结果 async for response in self._execute_coze_request(): yield response + except (ImagePayloadTooLargeError, MemoryError, OSError): + # Keep resource failures visible to the caller; do not turn them into + # a normal agent response that could trigger an unrelated retry. + self._transition_state(AgentState.ERROR) + raise except Exception as e: logger.error(f"Coze 请求失败:{str(e)}") self._transition_state(AgentState.ERROR) @@ -148,6 +166,21 @@ async def _execute_coze_request(self): # 处理历史上下文 if not self.auto_save_history and contexts: + image_options = ( + self.req.image_preparation_options or ImagePreparationOptions() + ) + image_limit = image_options.max_encoded_bytes + if image_limit is None: + image_limit = ImagePreparationOptions().max_encoded_bytes + validate_context_image_bytes( + contexts, + image_limit, + ) + contexts = await materialize_image_media_refs( + contexts, + ImageMediaStore(Path(get_astrbot_data_path()) / "media"), + ) + history_image_options = replace(image_options, enabled=False) for ctx in contexts: if is_checkpoint_message(ctx): continue @@ -169,7 +202,9 @@ async def _execute_coze_request(self): if url: file_id = ( await self._download_and_upload_image( - url, session_id + url, + session_id, + image_options=history_image_options, ) ) processed_content.append( @@ -179,6 +214,12 @@ async def _execute_coze_request(self): "file_url": url, } ) + except ( + ImagePayloadTooLargeError, + MemoryError, + OSError, + ): + raise except Exception as e: logger.warning(f"处理上下文图片失败: {e}") continue @@ -221,6 +262,8 @@ async def _execute_coze_request(self): "file_id": file_id, } ) + except (ImagePayloadTooLargeError, MemoryError, OSError): + raise except Exception as e: logger.warning( "处理图片失败 %s: %s", @@ -335,12 +378,19 @@ async def _download_and_upload_image( self, image_url: str, session_id: str | None = None, + *, + image_options: ImagePreparationOptions | None = None, ) -> str: """下载图片并上传到 Coze,返回 file_id""" import hashlib - # 计算哈希实现缓存 - cache_key = hashlib.md5(image_url.encode("utf-8")).hexdigest() + image_options = image_options or getattr( + getattr(self, "req", None), "image_preparation_options", None + ) + # Include preparation options so a changed byte budget cannot reuse an old upload. + cache_key = hashlib.sha256( + f"{image_url}\0{image_options!r}".encode() + ).hexdigest() if session_id: if session_id not in self.file_id_cache: @@ -352,10 +402,15 @@ async def _download_and_upload_image( return file_id try: - image_bytes = await MediaResolver( + image_data = await resolve_media_ref_to_base64_data( image_url, media_type="image", - ).to_bytes() + strict=True, + image_options=image_options, + ) + if image_data is None: + raise ValueError("Image preprocessing returned no data") + image_bytes = image_data.to_bytes() file_id = await self.api_client.upload_file(image_bytes) if session_id: @@ -364,6 +419,8 @@ async def _download_and_upload_image( return file_id + except (ImagePayloadTooLargeError, MemoryError, OSError): + raise except Exception as e: logger.error("处理图片失败 %s: %s", describe_media_ref(image_url), e) raise Exception(f"处理图片失败: {e!s}") from e diff --git a/astrbot/core/agent/runners/deerflow/deerflow_agent_runner.py b/astrbot/core/agent/runners/deerflow/deerflow_agent_runner.py index 4e3a40c5a0..40ab5d87ae 100644 --- a/astrbot/core/agent/runners/deerflow/deerflow_agent_runner.py +++ b/astrbot/core/agent/runners/deerflow/deerflow_agent_runner.py @@ -16,6 +16,10 @@ ProviderRequest, ) from astrbot.core.utils.config_number import coerce_int_config +from astrbot.core.utils.media_utils import ( + ImagePayloadTooLargeError, + ImagePreparationOptions, +) from ...hooks import BaseAgentRunHooks from ...response import AgentResponseData @@ -298,6 +302,11 @@ async def step(self): except asyncio.CancelledError: # Let caller manage cancellation semantics. raise + except (ImagePayloadTooLargeError, MemoryError, OSError): + # Keep resource failures visible to the caller; do not turn them into + # a normal agent response that could trigger an unrelated retry. + self._transition_state(AgentState.ERROR) + raise except Exception as e: err_msg = self._format_exception(e) logger.error(f"DeerFlow request failed: {err_msg}", exc_info=True) @@ -416,6 +425,7 @@ async def _build_messages_resolved( prompt: str, image_urls: list[str], system_prompt: str | None, + image_options: ImagePreparationOptions | None = None, ) -> list[dict[str, T.Any]]: """Build DeerFlow messages after materializing image references. @@ -423,6 +433,7 @@ async def _build_messages_resolved( prompt: User prompt text. image_urls: Image references accepted by MediaResolver. system_prompt: Optional system prompt prepended to the request. + image_options: Preparation limits for the active provider request. Returns: Messages payload for DeerFlow. @@ -434,7 +445,11 @@ async def _build_messages_resolved( messages.append( { "role": "user", - "content": await build_user_content_resolved(prompt, image_urls), + "content": await build_user_content_resolved( + prompt, + image_urls, + image_options=image_options, + ), }, ) return messages @@ -497,6 +512,9 @@ async def _build_payload_resolved( """ runtime_configurable = self._build_runtime_configurable(thread_id) + image_options = getattr( + getattr(self, "req", None), "image_preparation_options", None + ) return { "assistant_id": self.assistant_id, "input": { @@ -504,6 +522,7 @@ async def _build_payload_resolved( prompt, image_urls, system_prompt, + image_options=image_options, ), }, "stream_mode": ["values", "messages-tuple", "custom"], diff --git a/astrbot/core/agent/runners/deerflow/deerflow_content_mapper.py b/astrbot/core/agent/runners/deerflow/deerflow_content_mapper.py index 621fac226b..9ae19dbcac 100644 --- a/astrbot/core/agent/runners/deerflow/deerflow_content_mapper.py +++ b/astrbot/core/agent/runners/deerflow/deerflow_content_mapper.py @@ -6,6 +6,8 @@ from astrbot import logger from astrbot.core.message.message_event_result import MessageChain from astrbot.core.utils.media_utils import ( + ImagePayloadTooLargeError, + ImagePreparationOptions, describe_media_ref, resolve_media_ref_to_base64_data, ) @@ -98,13 +100,18 @@ def build_user_content(prompt: str, image_urls: list[str]) -> Any: return content -async def build_user_content_resolved(prompt: str, image_urls: list[str]) -> Any: +async def build_user_content_resolved( + prompt: str, + image_urls: list[str], + image_options: ImagePreparationOptions | None = None, +) -> Any: """Build DeerFlow user content after resolving all supported image refs. Args: prompt: User text to include before image blocks. image_urls: Image references from plugins or message attachments. Supports local paths, HTTP(S), file URIs, base64://, data URIs, and bare base64. + image_options: Preparation limits for the active provider request. Returns: Plain text when no images are present; otherwise a multimodal content list. @@ -136,7 +143,10 @@ async def build_user_content_resolved(prompt: str, image_urls: list[str]) -> Any image_data = await resolve_media_ref_to_base64_data( image_ref, media_type="image", + image_options=image_options, ) + except (ImagePayloadTooLargeError, MemoryError, OSError): + raise except Exception as exc: skipped_invalid_images += 1 logger.debug( diff --git a/astrbot/core/agent/runners/dify/dify_agent_runner.py b/astrbot/core/agent/runners/dify/dify_agent_runner.py index c2875c70a4..cfe7a614f5 100644 --- a/astrbot/core/agent/runners/dify/dify_agent_runner.py +++ b/astrbot/core/agent/runners/dify/dify_agent_runner.py @@ -8,7 +8,11 @@ LLMResponse, ProviderRequest, ) -from astrbot.core.utils.media_utils import MediaResolver +from astrbot.core.utils.media_utils import ( + ImagePayloadTooLargeError, + MediaResolver, + resolve_media_ref_to_base64_data, +) from ...hooks import BaseAgentRunHooks from ...response import AgentResponseData @@ -80,6 +84,11 @@ async def step(self): # 执行 Dify 请求并处理结果 async for response in self._execute_dify_request(): yield response + except (ImagePayloadTooLargeError, MemoryError, OSError): + # Keep resource failures visible to the caller; do not turn them into + # a normal agent response that could trigger an unrelated retry. + self._transition_state(AgentState.ERROR) + raise except Exception as e: logger.error(f"Dify 请求失败:{str(e)}") self._transition_state(AgentState.ERROR) @@ -108,10 +117,15 @@ async def _upload_image_for_dify( image_url: str, session_id: str, ) -> dict[str, str] | None: - image_data = await MediaResolver( + image_options = getattr( + getattr(self, "req", None), "image_preparation_options", None + ) + image_data = await resolve_media_ref_to_base64_data( image_url, media_type="image", - ).to_base64_data(strict=True) + strict=True, + image_options=image_options, + ) if image_data is None: logger.warning("Dify 图片预处理结果为空,将忽略。") return None @@ -159,6 +173,8 @@ async def _execute_dify_request(self): for image_url in image_urls: try: image_payload = await self._upload_image_for_dify(image_url, session_id) + except (ImagePayloadTooLargeError, MemoryError, OSError): + raise except Exception as e: logger.warning(f"上传图片失败:{e}") continue diff --git a/astrbot/core/agent/runners/tool_loop_agent_runner.py b/astrbot/core/agent/runners/tool_loop_agent_runner.py index c9787ed6f0..8927378249 100644 --- a/astrbot/core/agent/runners/tool_loop_agent_runner.py +++ b/astrbot/core/agent/runners/tool_loop_agent_runner.py @@ -28,7 +28,7 @@ from astrbot.core.agent.message import ImageURLPart, TextPart, ThinkPart from astrbot.core.agent.tool import FunctionTool, ToolSet from astrbot.core.agent.tool_image_cache import tool_image_cache -from astrbot.core.exceptions import EmptyModelOutputError +from astrbot.core.exceptions import EmptyModelOutputError, ProviderRequestTooLargeError from astrbot.core.message.components import Json from astrbot.core.message.message_event_result import ( MessageChain, @@ -46,9 +46,25 @@ sanitize_contexts_by_modalities, ) from astrbot.core.provider.provider import Provider +from astrbot.core.utils.astrbot_path import get_astrbot_data_path +from astrbot.core.utils.image_media_store import ( + ImageMediaStore, + materialize_image_media_refs, + persist_inline_image_refs, +) +from astrbot.core.utils.media_utils import ( + ImagePayloadTooLargeError, + ImagePreparationInput, + get_image_preparation_options, + prepare_image_source, +) from ..context.compressor import ContextCompressor from ..context.config import ContextConfig +from ..context.image_budget import ( + get_image_encoded_byte_limit, + validate_context_image_bytes, +) from ..context.manager import ContextManager from ..context.token_counter import EstimateTokenCounter, TokenCounter from ..hooks import BaseAgentRunHooks @@ -229,9 +245,13 @@ async def reset( request_max_retries: int | None = None, tool_result_overflow_dir: str | None = None, read_tool: FunctionTool | None = None, + image_media_store: ImageMediaStore | None = None, **kwargs: T.Any, ) -> None: self.req = request + self.image_media_store = image_media_store or ImageMediaStore( + Path(get_astrbot_data_path()) / "media" + ) self.streaming = streaming self.enforce_max_turns = enforce_max_turns self.llm_compress_instruction = llm_compress_instruction @@ -255,6 +275,7 @@ async def reset( llm_compress_provider=self.llm_compress_provider, custom_token_counter=self.custom_token_counter, custom_compressor=self.custom_compressor, + image_media_store=self.image_media_store, ) self.request_context_manager = ContextManager( self.request_context_manager_config @@ -314,6 +335,12 @@ async def reset( or request.extra_user_content_parts ): m = await self._assemble_request_context_for_provider(request) + if self.image_media_store is not None: + m = ( + await asyncio.to_thread( + persist_inline_image_refs, [m], self.image_media_store + ) + )[0] messages.append(Message.model_validate(m)) if request.system_prompt: messages.insert( @@ -334,6 +361,15 @@ async def _assemble_request_context_for_provider( self, request: ProviderRequest, ) -> dict[str, T.Any]: + request = replace( + request, + image_preparation_options=( + request.image_preparation_options + or get_image_preparation_options( + getattr(self.provider, "provider_settings", {}) + ) + ), + ) modalities = self.provider.provider_config.get("modalities", None) if not modalities: # Unconfigured (None or empty list) defaults to support all modalities for backward compatibility return await request.assemble_context() @@ -497,12 +533,41 @@ async def _await_or_stop( abort_task.cancel() await asyncio.gather(abort_task, return_exceptions=True) + async def _prepare_provider_contexts( + self, contexts: list[Message] | list[dict[str, T.Any]] + ) -> list[Message] | list[dict[str, T.Any]]: + """Prepare a request-local image view for normal and tool-requery calls. + + Args: + contexts: Selected history and any request-local instructions. + + Returns: + Provider-compatible messages without changing persisted history. + + Raises: + ImagePayloadTooLargeError: An image exceeds the encoded-byte limit. + MemoryError: Image materialization exhausts available memory. + """ + selected_contexts = self._sanitize_contexts_for_provider(contexts) + limit = get_image_encoded_byte_limit( + getattr(self.provider, "provider_settings", {}) + ) + validate_context_image_bytes(selected_contexts, limit) + return await materialize_image_media_refs( + selected_contexts, + self.image_media_store + or ImageMediaStore(Path(get_astrbot_data_path()) / "media"), + ) + async def _iter_llm_responses( self, *, include_model: bool = True ) -> T.AsyncGenerator[LLMResponse, None]: """Yields chunks *and* a final LLMResponse.""" + provider_contexts = await self._prepare_provider_contexts( + self.run_context.messages + ) payload = { - "contexts": self._sanitize_contexts_for_provider(self.run_context.messages), + "contexts": provider_contexts, "func_tool": self._func_tool_for_provider(), "session_id": self.req.session_id, "extra_user_content_parts": self.req.extra_user_content_parts, # list[ContentPart] @@ -622,7 +687,22 @@ async def _iter_llm_responses_with_fallback( raise if self._is_stop_requested(): return + except ( + MemoryError, + ImagePayloadTooLargeError, + ProviderRequestTooLargeError, + ): + raise except Exception as exc: # noqa: BLE001 + response = getattr(exc, "response", None) + status_code = getattr(exc, "status_code", None) or getattr( + response, "status_code", None + ) + if status_code == 413: + raise ProviderRequestTooLargeError( + "The provider rejected this request because it is too large " + "(HTTP 413). Reduce the active context or image size." + ) from exc last_exception = exc logger.warning( "Chat Model %s request error: %s", @@ -1060,11 +1140,29 @@ async def step(self): # Build user message with images for LLM to review image_parts = [] for cached_img in cached_images: - img_data = tool_image_cache.get_image_base64_by_path( - cached_img.file_path, cached_img.mime_type + img_data = await prepare_image_source( + ImagePreparationInput( + cached_img.file_path, + source_kind=( + "cua_screenshot" + if cached_img.tool_name == "astrbot_cua_screenshot" + else "tool_image" + ), + ), + options=replace( + self.req.image_preparation_options + or get_image_preparation_options( + getattr(self.provider, "provider_settings", {}) + ), + preserve_dimensions=cached_img.tool_name + == "astrbot_cua_screenshot", + ), ) if img_data: - base64_data, mime_type = img_data + base64_data, mime_type = ( + img_data.base64_data, + img_data.mime_type, + ) image_parts.append( TextPart( text=f"[Image from tool '{cached_img.tool_name}', path='{cached_img.file_path}']" @@ -1079,9 +1177,15 @@ async def step(self): ) ) if image_parts: - self.run_context.messages.append( - Message(role="user", content=image_parts) - ) + image_message = Message(role="user", content=image_parts) + if self.image_media_store is not None: + stored = await asyncio.to_thread( + persist_inline_image_refs, + [image_message.model_dump()], + self.image_media_store, + ) + image_message = Message.model_validate(stored[0]) + self.run_context.messages.append(image_message) logger.debug( f"Appended {len(cached_images)} cached image(s) to context for LLM review" ) @@ -1349,9 +1453,9 @@ def _append_tool_call_result(tool_call_id: str, content: str) -> None: ) except Exception as e: logger.error(f"Error in on_tool_end hook: {e}", exc_info=True) + except (MemoryError, _ToolExecutionInterrupted): + raise except Exception as e: - if isinstance(e, _ToolExecutionInterrupted): - raise logger.warning(traceback.format_exc()) _append_tool_call_result( func_tool_id, @@ -1443,7 +1547,7 @@ async def _resolve_tool_exec( contexts = self._build_tool_requery_context(tool_names) requery_resp = await self._await_or_stop( self.provider.text_chat( - contexts=self._sanitize_contexts_for_provider(contexts), + contexts=await self._prepare_provider_contexts(contexts), func_tool=param_subset, model=self.req.model, session_id=self.req.session_id, @@ -1473,7 +1577,7 @@ async def _resolve_tool_exec( ) repair_resp = await self._await_or_stop( self.provider.text_chat( - contexts=self._sanitize_contexts_for_provider( + contexts=await self._prepare_provider_contexts( repair_contexts ), func_tool=param_subset, diff --git a/astrbot/core/astr_main_agent.py b/astrbot/core/astr_main_agent.py index ce86e2c18e..c74494754e 100644 --- a/astrbot/core/astr_main_agent.py +++ b/astrbot/core/astr_main_agent.py @@ -32,6 +32,7 @@ from astrbot.core.computer.booters.local import resolve_windows_shell from astrbot.core.conversation_mgr import Conversation from astrbot.core.db import BaseDatabase +from astrbot.core.exceptions import ProviderRequestTooLargeError from astrbot.core.message.components import File, Image, Record, Reply, Video from astrbot.core.persona_error_reply import ( extract_persona_custom_error_message_from_persona, @@ -108,7 +109,12 @@ ) from astrbot.core.utils.file_extract import extract_file_moonshotai from astrbot.core.utils.llm_metadata import LLM_METADATAS -from astrbot.core.utils.media_utils import is_file_uri, is_recoverable_image_error +from astrbot.core.utils.media_utils import ( + ImagePayloadTooLargeError, + get_image_preparation_options, + is_file_uri, + is_recoverable_image_error, +) from astrbot.core.utils.quoted_message.settings import ( SETTINGS as DEFAULT_QUOTED_MESSAGE_SETTINGS, ) @@ -726,7 +732,7 @@ async def _ensure_img_caption( caption = await _request_img_caption( image_caption_provider, cfg, - req.image_urls, + list(req.image_urls), plugin_context, ) if caption: @@ -734,6 +740,8 @@ async def _ensure_img_caption( TextPart(text=f"{caption}") ) req.image_urls = [] + except (ImagePayloadTooLargeError, MemoryError, ProviderRequestTooLargeError): + raise except Exception as exc: # noqa: BLE001 logger.error("处理图片描述失败: %s", exc) req.extra_user_content_parts.append(TextPart(text="[Image Captioning Failed]")) @@ -888,8 +896,14 @@ async def _process_quote_message( ) else: logger.warning("No provider found for image captioning in quote.") - except Exception as exc: - logger.error("Quote image captioning failed (%s).", type(exc).__name__) + except ( + ImagePayloadTooLargeError, + MemoryError, + ProviderRequestTooLargeError, + ): + raise + except BaseException as exc: + logger.error("处理引用图片失败: %s", exc) quoted_content = "\n".join(content_parts) quoted_text = f"\n{quoted_content}\n" @@ -1730,6 +1744,11 @@ async def build_main_agent( _apply_web_search_citation_prompt(event, req) + if req.image_preparation_options is None: + req.image_preparation_options = get_image_preparation_options( + config.provider_settings + ) + reset_coro = agent_runner.reset( provider=provider, request=req, diff --git a/astrbot/core/computer/file_read_utils.py b/astrbot/core/computer/file_read_utils.py index 0e6b2fce1a..55a48c0f6b 100644 --- a/astrbot/core/computer/file_read_utils.py +++ b/astrbot/core/computer/file_read_utils.py @@ -16,13 +16,6 @@ from astrbot.core.agent.context.token_counter import EstimateTokenCounter from astrbot.core.agent.message import Message from astrbot.core.agent.tool import ToolExecResult -from astrbot.core.utils.astrbot_path import get_astrbot_temp_path -from astrbot.core.utils.media_utils import ( - IMAGE_COMPRESS_DEFAULT_MAX_SIZE, - IMAGE_COMPRESS_DEFAULT_OPTIMIZE, - IMAGE_COMPRESS_DEFAULT_QUALITY, - _compress_image_sync, -) from .booters.base import ComputerBooter from .local_file_security import open_file_in_allowed_roots @@ -309,25 +302,6 @@ def _run() -> dict[str, str | int]: return await to_thread(_run) -async def _read_local_image_base64( - path: str, - file_descriptor: int | None = None, -) -> dict[str, str | int]: - def _run() -> dict[str, str | int]: - if file_descriptor is None: - data = Path(path).read_bytes() - else: - with os.fdopen(os.dup(file_descriptor), "rb") as file_obj: - file_obj.seek(0) - data = file_obj.read() - return { - "size_bytes": len(data), - "base64": base64.b64encode(data).decode("utf-8"), - } - - return await to_thread(_run) - - async def _read_local_file_bytes( path: str, file_descriptor: int | None = None, @@ -343,33 +317,6 @@ def _run() -> bytes: return await to_thread(_run) -async def _compress_image_bytes_to_base64(data: bytes) -> dict[str, str | int]: - def _run() -> dict[str, str | int]: - temp_dir = Path(get_astrbot_temp_path()) - temp_dir.mkdir(parents=True, exist_ok=True) - compressed_path = Path( - _compress_image_sync( - data, - temp_dir, - IMAGE_COMPRESS_DEFAULT_MAX_SIZE, - IMAGE_COMPRESS_DEFAULT_QUALITY, - IMAGE_COMPRESS_DEFAULT_OPTIMIZE, - ) - ) - try: - compressed_bytes = compressed_path.read_bytes() - finally: - compressed_path.unlink(missing_ok=True) - - return { - "size_bytes": len(compressed_bytes), - "base64": base64.b64encode(compressed_bytes).decode("utf-8"), - "mime_type": "image/jpeg", - } - - return await to_thread(_run) - - def _detect_image_mime(sample: bytes) -> str | None: if sample.startswith(b"\x89PNG\r\n\x1a\n"): return "image/png" @@ -757,32 +704,29 @@ async def read_file_tool_result( if probe.kind == "image": if local_mode: - image_payload = await _read_local_image_base64( - path, - local_file_descriptor, - ) + try: + raw_bytes = await _read_local_file_bytes(path, local_file_descriptor) + except OSError as exc: + return f"Error reading file: failed to read image: {exc}" + image_base64_data = base64.b64encode(raw_bytes).decode("utf-8") + mime_type = probe.mime_type or "image/jpeg" else: image_payload = await _exec_python_json( booter, _build_image_read_script(path), action="image read", ) - raw_base64_data = str(image_payload.get("base64", "") or "") - if not raw_base64_data: - return "Error reading file: image payload is empty." - raw_bytes = base64.b64decode(raw_base64_data) - compressed_payload = await _compress_image_bytes_to_base64(raw_bytes) - compressed_base64_data = str(compressed_payload.get("base64", "") or "") - if not compressed_base64_data: - return "Error reading file: compressed image payload is empty." + raw_base64_data = str(image_payload.get("base64", "") or "") + if not raw_base64_data: + return "Error reading file: image payload is empty." + image_base64_data = raw_base64_data + mime_type = probe.mime_type or "image/jpeg" return mcp.types.CallToolResult( content=[ mcp.types.ImageContent( type="image", - data=compressed_base64_data, - mimeType=str( - compressed_payload.get("mime_type", "") or "image/jpeg" - ), + data=image_base64_data, + mimeType=mime_type, ) ] ) diff --git a/astrbot/core/config/default.py b/astrbot/core/config/default.py index f8c46c26e4..862af86a72 100644 --- a/astrbot/core/config/default.py +++ b/astrbot/core/config/default.py @@ -202,6 +202,7 @@ def get_local_permission_defaults(system: str | None = None) -> dict: "image_compress_options": { "max_size": 1280, "quality": 95, + "max_encoded_bytes": 4194304, }, }, "agent_runner": { @@ -4232,6 +4233,15 @@ def get_local_permission_defaults(system: str | None = None) -> dict: }, "slider": {"min": 1, "max": 100, "step": 1}, }, + "provider_settings.image_compress_options.max_encoded_bytes": { + "description": "最大编码大小", + "type": "int", + "hint": "压缩后单张图片的最大 Base64 编码大小,单位为字节。超过限制且无法进一步压缩时会报错。", + "condition": { + "provider_settings.image_compress_enabled": True, + }, + "slider": {"min": 262144, "max": 16777216, "step": 262144}, + }, "provider_settings.prompt_prefix": { "description": "用户提示词", "type": "string", diff --git a/astrbot/core/exceptions.py b/astrbot/core/exceptions.py index 76c71af573..9915ed3c05 100644 --- a/astrbot/core/exceptions.py +++ b/astrbot/core/exceptions.py @@ -13,6 +13,10 @@ class EmptyModelOutputError(AstrBotError): """Raised when the model response contains no usable assistant output.""" +class ProviderRequestTooLargeError(AstrBotError): + """Raised when a provider rejects the serialized request size.""" + + class KnowledgeBaseUploadError(AstrBotError): """Raised when knowledge base upload fails with a user-facing message.""" diff --git a/astrbot/core/pipeline/process_stage/method/agent_sub_stages/internal.py b/astrbot/core/pipeline/process_stage/method/agent_sub_stages/internal.py index 322506e3a8..21bf1a49b0 100644 --- a/astrbot/core/pipeline/process_stage/method/agent_sub_stages/internal.py +++ b/astrbot/core/pipeline/process_stage/method/agent_sub_stages/internal.py @@ -538,12 +538,14 @@ async def process( unregister_active_runner(event.unified_msg_origin, agent_runner) except Exception as e: - logger.error(f"Error occurred while processing agent: {e}") + logger.error( + f"Error occurred while processing agent: {type(e).__name__}: {e}" + ) custom_error_message = extract_persona_custom_error_message_from_event( event ) error_text = custom_error_message or ( - f"Error occurred while processing agent request: {e}" + f"Error occurred while processing agent request: {type(e).__name__}: {e}" ) await event.send(MessageChain().message(error_text)) finally: diff --git a/astrbot/core/pipeline/process_stage/method/agent_sub_stages/third_party.py b/astrbot/core/pipeline/process_stage/method/agent_sub_stages/third_party.py index 93836c1fe8..1cd87f41e4 100644 --- a/astrbot/core/pipeline/process_stage/method/agent_sub_stages/third_party.py +++ b/astrbot/core/pipeline/process_stage/method/agent_sub_stages/third_party.py @@ -36,6 +36,7 @@ ) from astrbot.core.star.star_handler import EventType from astrbot.core.utils.config_number import coerce_int_config +from astrbot.core.utils.media_utils import get_image_preparation_options from astrbot.core.utils.metrics import Metric from .....astr_agent_context import AgentContextWrapper, AstrAgentContext @@ -289,10 +290,21 @@ async def process( req = ProviderRequest() req.session_id = event.unified_msg_origin req.prompt = event.message_str[len(provider_wake_prefix) :] + req.image_preparation_options = get_image_preparation_options( + self.conf.get("provider_settings") + ) for comp in event.message_obj.message: if isinstance(comp, Image): - image_path = await comp.convert_to_base64() - req.image_urls.append(image_path) + image_ref = ( + getattr(comp, "path", None) + or getattr(comp, "url", None) + or getattr(comp, "file", None) + ) + if not isinstance(image_ref, str): + image_ref = await comp.convert_to_file_path() + if not image_ref: + raise ValueError("Image attachment has no usable reference") + req.image_urls.append(image_ref) elif isinstance(comp, Record): audio_path = await comp.convert_to_file_path() req.audio_urls.append(audio_path) diff --git a/astrbot/core/provider/entities.py b/astrbot/core/provider/entities.py index 382932b857..b7ebd95d47 100644 --- a/astrbot/core/provider/entities.py +++ b/astrbot/core/provider/entities.py @@ -24,9 +24,12 @@ from astrbot.core.db.po import Conversation from astrbot.core.message.message_event_result import MessageChain from astrbot.core.utils.media_utils import ( + ImagePayloadTooLargeError, + ImagePreparationInput, + ImagePreparationOptions, MediaResolver, is_recoverable_image_error, - resolve_image_ref_to_base64_data, + prepare_image_source, ) @@ -118,6 +121,8 @@ class ProviderRequest: """附加的上次请求后工具调用的结果。参考: https://platform.openai.com/docs/guides/function-calling#handling-function-calls""" model: str | None = None """模型名称,为 None 时使用提供商的默认模型""" + image_preparation_options: ImagePreparationOptions | None = None + """当前请求的图片准备配置,由上层 Agent 配置注入。""" def __repr__(self) -> str: return ( @@ -215,14 +220,20 @@ async def assemble_context(self) -> dict: dumped = ( part if isinstance(part, dict) else part.model_dump_for_context() ) - # Capture bytes before event cleanup. Extra image paths must also - # reach providers and persisted history as portable data URIs. if isinstance(dumped, dict) and dumped.get("type") == "image_url": image_url = dumped.get("image_url") url = image_url.get("url") if isinstance(image_url, dict) else None if isinstance(url, str) and url: try: - resolved = await resolve_image_ref_to_base64_data(url) + prepared = await prepare_image_source( + ImagePreparationInput( + url, + source_kind="extra_user_content", + ), + options=self.image_preparation_options, + ) + except (ImagePayloadTooLargeError, MemoryError): + raise except Exception as exc: if not is_recoverable_image_error(exc): raise @@ -232,26 +243,25 @@ async def assemble_context(self) -> dict: ) image_capture_failed = True continue - if resolved is None: - logger.warning( - "Image source capture returned no data; skipping image." - ) - image_capture_failed = True - continue dumped = { **dumped, "image_url": { **image_url, - "url": resolved.to_data_url(), + "url": prepared.to_data_url(), }, } content_blocks.append(dumped) - # 3. Read image references without resizing or transcoding. + # 3. 图片内容 if self.image_urls: for image_url in self.image_urls: try: - image_data = await resolve_image_ref_to_base64_data(image_url) + image_data = await prepare_image_source( + ImagePreparationInput(image_url, source_kind="request_image"), + options=self.image_preparation_options, + ) + except (ImagePayloadTooLargeError, MemoryError): + raise except Exception as exc: if not is_recoverable_image_error(exc): raise @@ -261,12 +271,6 @@ async def assemble_context(self) -> dict: ) image_capture_failed = True continue - if image_data is None: - logger.warning( - "Image source capture returned no data; skipping image." - ) - image_capture_failed = True - continue content_blocks.append( { "type": "image_url", diff --git a/astrbot/core/provider/modalities.py b/astrbot/core/provider/modalities.py index 66ac74e9b7..b14142760a 100644 --- a/astrbot/core/provider/modalities.py +++ b/astrbot/core/provider/modalities.py @@ -91,7 +91,11 @@ def sanitize_contexts_by_modalities( for part in content: if isinstance(part, dict): part_type = str(part.get("type", "")).lower() - if not supports_image and part_type in {"image_url", "image"}: + if not supports_image and part_type in { + "image_url", + "image", + "image_media_ref", + }: removed_any_multimodal = True stats.fixed_image_blocks += 1 filtered_parts.append({"type": "text", "text": "[Image]"}) @@ -132,7 +136,7 @@ def _tool_result_placeholder(content: Any) -> str: part_type = str(part.get("type", "")).lower() if part_type == "text": text_parts.append(str(part.get("text", ""))) - elif part_type in {"image_url", "image"}: + elif part_type in {"image_url", "image", "image_media_ref"}: text_parts.append("[Image]") elif part_type in {"audio_url", "input_audio"}: text_parts.append("[Audio]") diff --git a/astrbot/core/provider/sources/anthropic_source.py b/astrbot/core/provider/sources/anthropic_source.py index 5a21737b44..03f33867ee 100644 --- a/astrbot/core/provider/sources/anthropic_source.py +++ b/astrbot/core/provider/sources/anthropic_source.py @@ -1,6 +1,7 @@ import base64 import json from collections.abc import AsyncGenerator +from pathlib import Path from typing import Any, Literal import anthropic @@ -12,12 +13,23 @@ from astrbot import logger from astrbot.api.provider import Provider +from astrbot.core.agent.context.image_budget import ( + get_image_encoded_byte_limit, + validate_context_image_bytes, +) from astrbot.core.agent.message import AudioURLPart, ContentPart, ImageURLPart, TextPart from astrbot.core.exceptions import EmptyModelOutputError from astrbot.core.provider.entities import LLMResponse, TokenUsage from astrbot.core.provider.func_tool_manager import ToolSet +from astrbot.core.provider.modalities import sanitize_contexts_by_modalities +from astrbot.core.utils.astrbot_path import get_astrbot_data_path +from astrbot.core.utils.image_media_store import ( + ImageMediaStore, + materialize_image_media_refs, +) from astrbot.core.utils.media_utils import ( describe_media_ref, + get_image_preparation_options, resolve_media_ref_to_base64_data, ) from astrbot.core.utils.network_utils import ( @@ -274,11 +286,14 @@ def _prepare_payload(self, messages: list[dict]): if url.startswith("data:"): try: _, base64_data = url.split(",", 1) - # Detect actual image format from binary data - image_bytes = base64.b64decode(base64_data) - media_type = self._detect_image_mime_type( - image_bytes + # Decode only the header-sized prefix. The + # complete payload remains the provider's + # wire data and must not be copied merely + # to detect its MIME type. + prefix = base64.b64decode( + "".join(base64_data.split())[:64] ) + media_type = self._detect_image_mime_type(prefix) converted_content.append( { "type": "image", @@ -780,6 +795,15 @@ async def text_chat( ) -> LLMResponse: if contexts is None: contexts = [] + contexts, _ = sanitize_contexts_by_modalities( + contexts, self.provider_config.get("modalities") + ) + validate_context_image_bytes( + contexts, get_image_encoded_byte_limit(self.provider_settings) + ) + contexts = await materialize_image_media_refs( + contexts, ImageMediaStore(Path(get_astrbot_data_path()) / "media") + ) new_record = None if prompt is not None: new_record = await self.assemble_context( @@ -853,6 +877,15 @@ async def text_chat_stream( ): if contexts is None: contexts = [] + contexts, _ = sanitize_contexts_by_modalities( + contexts, self.provider_config.get("modalities") + ) + validate_context_image_bytes( + contexts, get_image_encoded_byte_limit(self.provider_settings) + ) + contexts = await materialize_image_media_refs( + contexts, ImageMediaStore(Path(get_astrbot_data_path()) / "media") + ) new_record = None if prompt is not None: new_record = await self.assemble_context( @@ -927,6 +960,7 @@ async def resolve_image_url(image_url: str) -> dict | None: image_data = await resolve_media_ref_to_base64_data( image_url, media_type="image", + image_options=get_image_preparation_options(self.provider_settings), ) if not image_data: logger.warning("图片预处理结果为空,将忽略。") @@ -999,6 +1033,7 @@ async def encode_image_bs64(self, image_url: str) -> tuple[str, str]: image_url, media_type="image", strict=True, + image_options=get_image_preparation_options(self.provider_settings), ) if image_data is None: raise RuntimeError( diff --git a/astrbot/core/provider/sources/gemini_source.py b/astrbot/core/provider/sources/gemini_source.py index 00268b2c0b..54eb484138 100644 --- a/astrbot/core/provider/sources/gemini_source.py +++ b/astrbot/core/provider/sources/gemini_source.py @@ -4,6 +4,8 @@ import logging import random from collections.abc import AsyncGenerator +from dataclasses import replace +from pathlib import Path from typing import Literal, cast import httpx @@ -14,13 +16,25 @@ import astrbot.core.message.components as Comp from astrbot import logger from astrbot.api.provider import Provider +from astrbot.core.agent.context.image_budget import ( + get_image_encoded_byte_limit, + validate_context_image_bytes, +) from astrbot.core.agent.message import AudioURLPart, ContentPart, ImageURLPart, TextPart -from astrbot.core.exceptions import EmptyModelOutputError +from astrbot.core.exceptions import EmptyModelOutputError, ProviderRequestTooLargeError from astrbot.core.message.message_event_result import MessageChain from astrbot.core.provider.entities import LLMResponse, TokenUsage from astrbot.core.provider.func_tool_manager import ToolSet +from astrbot.core.provider.modalities import sanitize_contexts_by_modalities +from astrbot.core.utils.astrbot_path import get_astrbot_data_path +from astrbot.core.utils.image_media_store import ( + ImageMediaStore, + materialize_image_media_refs, +) from astrbot.core.utils.media_utils import ( describe_media_ref, + detect_image_mime_type, + get_image_preparation_options, resolve_media_ref_to_base64_data, ) from astrbot.core.utils.network_utils import is_connection_error, log_connection_failure @@ -133,6 +147,10 @@ def _init_safety_settings(self) -> None: async def _handle_api_error(self, e: APIError, keys: list[str]) -> bool: """处理API错误,返回是否需要重试""" + if getattr(e, "code", None) == 413: + raise ProviderRequestTooLargeError( + "The provider rejected the request because its serialized size is too large (HTTP 413)." + ) from e if e.message is None: e.message = "" @@ -299,6 +317,20 @@ async def _prepare_query_config( async def _prepare_conversation(self, payloads: dict) -> list[types.Content]: """准备 Gemini SDK 的 Content 列表""" + payloads = dict(payloads) + provider_config = getattr(self, "provider_config", {}) + provider_settings = getattr(self, "provider_settings", {}) + messages, _ = sanitize_contexts_by_modalities( + payloads.get("messages", []), provider_config.get("modalities") + ) + validate_context_image_bytes( + messages, get_image_encoded_byte_limit(provider_settings) + ) + payloads["messages"] = messages + payloads["messages"] = await materialize_image_media_refs( + payloads.get("messages", []), + ImageMediaStore(Path(get_astrbot_data_path()) / "media"), + ) def create_text_part(text: str) -> types.Part: content_a = text if text else " " @@ -308,10 +340,30 @@ def create_text_part(text: str) -> types.Part: async def process_image_url(image_url_dict: dict) -> types.Part: url = image_url_dict["url"] + if url.startswith("data:"): + header, separator, encoded = url.partition(",") + if separator and ";base64" in header.lower(): + image_bytes = base64.b64decode(encoded) + declared_mime = header[5:].split(";", 1)[0].strip() + detected_mime = await asyncio.to_thread( + detect_image_mime_type, + image_bytes, + default_mime_type=None, + ) + mime_type = detected_mime or declared_mime + if mime_type: + return types.Part.from_bytes( + data=image_bytes, + mime_type=mime_type, + ) image_data = await resolve_media_ref_to_base64_data( url, media_type="image", strict=True, + image_options=replace( + get_image_preparation_options(provider_settings), + enabled=False, + ), ) if image_data is None: raise ValueError( @@ -1007,6 +1059,7 @@ async def resolve_image_part(image_url: str) -> dict | None: image_data = await resolve_media_ref_to_base64_data( image_url, media_type="image", + image_options=get_image_preparation_options(self.provider_settings), ) if not image_data: logger.warning("Image preprocessing returned no data; ignoring it.") @@ -1103,6 +1156,7 @@ async def encode_image_bs64(self, image_url: str) -> str: image_url, media_type="image", strict=True, + image_options=get_image_preparation_options(self.provider_settings), ) if image_data is None: raise RuntimeError( diff --git a/astrbot/core/provider/sources/openai_responses_source.py b/astrbot/core/provider/sources/openai_responses_source.py index c5cb9bdb82..76a19040e6 100644 --- a/astrbot/core/provider/sources/openai_responses_source.py +++ b/astrbot/core/provider/sources/openai_responses_source.py @@ -2,17 +2,28 @@ import inspect import json from collections.abc import AsyncGenerator +from pathlib import Path from typing import Any from openai.types.responses import Response import astrbot.core.message.components as Comp from astrbot import logger +from astrbot.core.agent.context.image_budget import ( + get_image_encoded_byte_limit, + validate_context_image_bytes, +) from astrbot.core.agent.message import ContentPart, Message from astrbot.core.agent.tool import ToolSet from astrbot.core.exceptions import EmptyModelOutputError from astrbot.core.message.message_event_result import MessageChain from astrbot.core.provider.entities import LLMResponse, TokenUsage, ToolCallsResult +from astrbot.core.provider.modalities import sanitize_contexts_by_modalities +from astrbot.core.utils.astrbot_path import get_astrbot_data_path +from astrbot.core.utils.image_media_store import ( + ImageMediaStore, + materialize_image_media_refs, +) from ..register import register_provider_adapter from .openai_source import ProviderOpenAIOfficial @@ -259,6 +270,17 @@ async def _prepare_chat_payload( Returns: The Responses payload and its chat-format source context. """ + contexts = contexts or [] + contexts, _ = sanitize_contexts_by_modalities( + contexts, self.provider_config.get("modalities") + ) + validate_context_image_bytes( + contexts, get_image_encoded_byte_limit(self.provider_settings) + ) + contexts = await materialize_image_media_refs( + contexts, + ImageMediaStore(Path(get_astrbot_data_path()) / "media"), + ) context_query = copy.deepcopy(self._ensure_message_to_dicts(contexts)) if prompt is not None: context_query.append( diff --git a/astrbot/core/provider/sources/openai_source.py b/astrbot/core/provider/sources/openai_source.py index 7d5a9ffd9e..3cd2bc9547 100644 --- a/astrbot/core/provider/sources/openai_source.py +++ b/astrbot/core/provider/sources/openai_source.py @@ -5,6 +5,8 @@ import random import re from collections.abc import AsyncGenerator +from dataclasses import replace +from pathlib import Path from typing import Any, Literal import httpx @@ -18,6 +20,10 @@ import astrbot.core.message.components as Comp from astrbot import logger from astrbot.api.provider import Provider +from astrbot.core.agent.context.image_budget import ( + get_image_encoded_byte_limit, + validate_context_image_bytes, +) from astrbot.core.agent.message import ( AudioURLPart, ContentPart, @@ -26,11 +32,20 @@ TextPart, ) from astrbot.core.agent.tool import ToolSet -from astrbot.core.exceptions import EmptyModelOutputError +from astrbot.core.exceptions import EmptyModelOutputError, ProviderRequestTooLargeError from astrbot.core.message.message_event_result import MessageChain from astrbot.core.provider.entities import LLMResponse, TokenUsage, ToolCallsResult +from astrbot.core.provider.modalities import sanitize_contexts_by_modalities +from astrbot.core.utils.astrbot_path import get_astrbot_data_path +from astrbot.core.utils.image_media_store import ( + ImageMediaStore, + materialize_image_media_refs, +) from astrbot.core.utils.media_utils import ( + ImagePayloadTooLargeError, + ImagePreparationOptions, describe_media_ref, + get_image_preparation_options, resolve_media_ref_to_base64_data, ) from astrbot.core.utils.network_utils import ( @@ -181,11 +196,14 @@ async def _image_ref_to_data_url( image_ref: str, *, mode: Literal["safe", "strict"] = "safe", + options: ImagePreparationOptions | None = None, ) -> str | None: image_data = await resolve_media_ref_to_base64_data( image_ref, media_type="image", strict=mode == "strict", + image_options=options + or get_image_preparation_options(self.provider_settings), ) return image_data.to_data_url() if image_data else None @@ -194,8 +212,13 @@ async def _resolve_image_part( image_url: str, *, image_detail: str | None = None, + options: ImagePreparationOptions | None = None, ) -> dict | None: - image_data = await self._image_ref_to_data_url(image_url, mode="safe") + image_data = await self._image_ref_to_data_url( + image_url, + mode="safe", + options=options, + ) if not image_data: logger.warning("图片预处理结果为空,将忽略。") return None @@ -266,7 +289,12 @@ async def _resolve_audio_part(self, audio_ref: str) -> dict | None: }, } - async def _transform_content_part(self, part: dict) -> dict: + async def _transform_content_part( + self, + part: dict, + *, + options: ImagePreparationOptions | None = None, + ) -> dict: if not isinstance(part, dict): return part @@ -275,10 +303,20 @@ async def _transform_content_part(self, part: dict) -> dict: if not url: return part + # ProviderRequest.assemble_context() already materializes current + # images into a data URL. Keep that request-local representation + # unchanged so the adapter does not decode and encode it again. + if url.startswith("data:image/"): + return part + try: resolved_part = await self._resolve_image_part( - url, image_detail=image_detail + url, + image_detail=image_detail, + options=options, ) + except (ImagePayloadTooLargeError, MemoryError): + raise except Exception as exc: logger.warning( "图片 %s 预处理失败,将保留原始内容。错误: %s", @@ -298,19 +336,30 @@ async def _transform_content_part(self, part: dict) -> dict: return part - async def _materialize_message_image_parts(self, message: dict) -> dict: + async def _materialize_message_image_parts( + self, + message: dict, + *, + options: ImagePreparationOptions | None = None, + ) -> dict: content = message.get("content") if not isinstance(content, list): return {**message} - new_content = [await self._transform_content_part(part) for part in content] + new_content = [ + await self._transform_content_part(part, options=options) + for part in content + ] return {**message, "content": new_content} async def _materialize_context_image_parts( self, context_query: list[dict] ) -> list[dict]: + options = replace( + get_image_preparation_options(self.provider_settings), enabled=False + ) return [ - await self._materialize_message_image_parts(message) + await self._materialize_message_image_parts(message, options=options) for message in context_query ] @@ -961,6 +1010,15 @@ async def _prepare_chat_payload( """准备聊天所需的有效载荷和上下文""" if contexts is None: contexts = [] + contexts, _ = sanitize_contexts_by_modalities( + contexts, self.provider_config.get("modalities") + ) + validate_context_image_bytes( + contexts, get_image_encoded_byte_limit(self.provider_settings) + ) + contexts = await materialize_image_media_refs( + contexts, ImageMediaStore(Path(get_astrbot_data_path()) / "media") + ) new_record = None if prompt is not None: new_record = await self.assemble_context( @@ -1081,6 +1139,16 @@ async def _handle_api_error( image_fallback_used: bool = False, ) -> tuple: """处理API错误并尝试恢复""" + if isinstance(e, MemoryError): + raise e + status_code = getattr(e, "status_code", None) + response = getattr(e, "response", None) + if status_code is None and response is not None: + status_code = getattr(response, "status_code", None) + if status_code == 413: + raise ProviderRequestTooLargeError( + "The provider rejected the request because its serialized size is too large (HTTP 413)." + ) from e if "429" in str(e): logger.warning( f"API 调用过于频繁,尝试使用其他 Key 重试。当前 Key: {chosen_key[:12]}", diff --git a/astrbot/core/provider/sources/request_retry.py b/astrbot/core/provider/sources/request_retry.py index 14dd57d7ab..3123ca2d18 100644 --- a/astrbot/core/provider/sources/request_retry.py +++ b/astrbot/core/provider/sources/request_retry.py @@ -11,6 +11,7 @@ ) from astrbot import logger +from astrbot.core.exceptions import ProviderRequestTooLargeError from astrbot.core.utils.config_number import coerce_int_config from astrbot.core.utils.network_utils import is_connection_error @@ -42,6 +43,8 @@ def _is_retryable_provider_request_error( *, retry_rate_limits: bool, ) -> bool: + if isinstance(error, (MemoryError, ProviderRequestTooLargeError)): + return False if is_connection_error(error): return True @@ -123,7 +126,16 @@ async def retry_provider_request( async for attempt in retrying: with attempt: - return await request_factory() + try: + return await request_factory() + except (MemoryError, ProviderRequestTooLargeError): + raise + except Exception as error: + if _get_status_code(error) == 413: + raise ProviderRequestTooLargeError( + "The provider rejected the request because its serialized size is too large (HTTP 413)." + ) from error + raise raise RuntimeError("Provider request retry loop exited unexpectedly.") diff --git a/astrbot/core/tools/computer_tools/fs.py b/astrbot/core/tools/computer_tools/fs.py index 906a2f0d8a..11527d1469 100644 --- a/astrbot/core/tools/computer_tools/fs.py +++ b/astrbot/core/tools/computer_tools/fs.py @@ -438,6 +438,8 @@ async def call( ) except PermissionError as exc: return f"Error: {exc}" + except MemoryError: + raise except Exception as exc: logger.error(f"Error reading file: {exc}") return f"Error reading file: {exc}" diff --git a/astrbot/core/utils/image_media_store.py b/astrbot/core/utils/image_media_store.py new file mode 100644 index 0000000000..552a3496c1 --- /dev/null +++ b/astrbot/core/utils/image_media_store.py @@ -0,0 +1,357 @@ +"""Durable, content-addressed storage for conversation image bytes.""" + +from __future__ import annotations + +import base64 +import hashlib +import io +import json +import os +import tempfile +from dataclasses import asdict, dataclass +from pathlib import Path + +from PIL import Image + + +@dataclass(frozen=True, slots=True) +class ImageMediaRef: + """A versioned reference persisted in a conversation message. + + Args: + media_id: SHA-256 digest of the exact stored bytes. + mime_type: MIME type sent to a provider. + width: Pixel width, if the bytes are a readable image. + height: Pixel height, if the bytes are a readable image. + byte_size: Exact stored byte count. + detail: Provider image-detail metadata. + """ + + media_id: str + mime_type: str + width: int | None + height: int | None + byte_size: int + detail: str | None = None + version: int = 1 + image_id: str | None = None + + def __post_init__(self) -> None: + if ( + self.version != 1 + or len(self.media_id) != 64 + or any(char not in "0123456789abcdef" for char in self.media_id) + or self.byte_size < 0 + or not self.mime_type.startswith("image/") + or (self.image_id is not None and not isinstance(self.image_id, str)) + ): + raise ValueError("Invalid durable image reference") + + @property + def uri(self) -> str: + """Return the internal URI; this URI is never sent to a provider.""" + return f"astrbot-media:v{self.version}:{self.media_id}" + + def model_dump(self) -> dict[str, object]: + """Return a stable JSON-compatible history representation.""" + return {"type": "image_media_ref", **asdict(self)} + + +class ImageMediaStore: + """Store immutable image bytes outside the temporary directory.""" + + def __init__(self, root: Path) -> None: + """Create a store rooted under the configured data directory. + + Args: + root: Dedicated durable media directory, not a client path. + """ + self.root = root + + def put( + self, + data: bytes, + mime_type: str | None = None, + detail: str | None = None, + image_id: str | None = None, + ) -> ImageMediaRef: + """Atomically persist bytes and return their deduplicated reference. + + Args: + data: Exact prepared image bytes. + mime_type: Optional MIME type; Pillow detection is preferred. + detail: Provider image-detail metadata. + + Returns: + A reference whose object and metadata have both been verified. + + Raises: + OSError: The object or metadata cannot be committed atomically. + ValueError: The stored bytes are not a readable image. + """ + media_id = hashlib.sha256(data).hexdigest() + self.root.mkdir(parents=True, exist_ok=True) + if self.root.is_symlink() or not self.root.is_dir(): + raise OSError("invalid durable media root") + object_path = self.root / f"{media_id}.bin" + metadata_path = self.root / f"{media_id}.json" + detected_mime = "application/octet-stream" + width: int | None = None + height: int | None = None + try: + with Image.open(io.BytesIO(data)) as image: + width, height = image.size + detected_mime = Image.MIME.get(image.format, detected_mime) + except MemoryError: + raise + except Exception as exc: # noqa: BLE001 + raise ValueError("media bytes are not a readable image") from exc + ref = ImageMediaRef( + media_id, + mime_type or detected_mime, + width, + height, + len(data), + detail, + image_id=image_id, + ) + canonical = { + "media_id": ref.media_id, + "mime_type": detected_mime, + "width": ref.width, + "height": ref.height, + "byte_size": ref.byte_size, + "version": ref.version, + } + if object_path.exists() or metadata_path.exists(): + # Verify shared bytes first. An interrupted metadata write can then + # be completed through the same atomic path as a new object. + if object_path.is_symlink() or not object_path.is_file(): + raise OSError("incomplete durable media object") + if hashlib.sha256(object_path.read_bytes()).hexdigest() != media_id: + raise OSError("media hash verification failed") + if metadata_path.is_symlink(): + raise OSError("incomplete durable media object") + if metadata_path.exists(): + if not metadata_path.is_file(): + raise OSError("incomplete durable media object") + try: + persisted = json.loads(metadata_path.read_text()) + except (OSError, json.JSONDecodeError) as exc: + raise OSError("durable media metadata verification failed") from exc + if persisted != canonical: + raise OSError("durable media metadata verification failed") + self.read(ref, {media_id}) + return ref + temporary_paths: list[Path] = [] + try: + with tempfile.NamedTemporaryFile(dir=self.root, delete=False) as output: + object_tmp = Path(output.name) + temporary_paths.append(object_tmp) + output.write(data) + output.flush() + os.fsync(output.fileno()) + try: + os.link(temporary_paths[0], object_path) + except FileExistsError: + pass + if hashlib.sha256(object_path.read_bytes()).hexdigest() != media_id: + raise OSError("media hash verification failed") + with tempfile.NamedTemporaryFile( + dir=self.root, + prefix=f"{media_id}.", + suffix=".json", + mode="w", + delete=False, + ) as metadata_file: + metadata_path_tmp = Path(metadata_file.name) + temporary_paths.append(metadata_path_tmp) + metadata_file.write(json.dumps(canonical, sort_keys=True) + "\n") + metadata_file.flush() + os.fsync(metadata_file.fileno()) + try: + os.link(metadata_path_tmp, metadata_path) + except FileExistsError: + pass + if json.loads(metadata_path.read_text()) != canonical: + raise OSError("durable media metadata verification failed") + finally: + for path in temporary_paths: + path.unlink(missing_ok=True) + return ref + + def read(self, ref: ImageMediaRef, allowed_media_ids: set[str]) -> bytes: + """Read a reference only when the caller already proved access. + + Args: + ref: Persisted reference, not a client-supplied filesystem path. + allowed_media_ids: IDs belonging to the authorized conversation. + + Returns: + The exact bytes committed for the reference. + + Raises: + PermissionError: The reference is outside the authorized history. + FileNotFoundError: The durable object is missing. + OSError: The object hash no longer matches its reference. + """ + if ref.media_id not in allowed_media_ids: + raise PermissionError("media reference is not authorized") + if self.root.is_symlink() or not self.root.is_dir(): + raise OSError("invalid durable media root") + path = self.root / f"{ref.media_id}.bin" + if path.is_symlink(): + raise OSError("invalid durable media object") + if not path.exists(): + raise FileNotFoundError("durable media object is unavailable") + if not path.is_file(): + raise OSError("invalid durable media object") + metadata_path = self.root / f"{ref.media_id}.json" + if metadata_path.is_symlink() or not metadata_path.exists(): + raise OSError("incomplete durable media object") + try: + metadata = json.loads(metadata_path.read_text()) + metadata.pop("type", None) + persisted_ref = ImageMediaRef(**metadata) + except (OSError, json.JSONDecodeError, TypeError, ValueError) as exc: + raise OSError("durable media metadata verification failed") from exc + if ( + persisted_ref.media_id, + persisted_ref.width, + persisted_ref.height, + persisted_ref.byte_size, + persisted_ref.version, + ) != ( + ref.media_id, + ref.width, + ref.height, + ref.byte_size, + ref.version, + ): + raise OSError("durable media metadata verification failed") + data = path.read_bytes() + if ( + len(data) != ref.byte_size + or hashlib.sha256(data).hexdigest() != ref.media_id + ): + raise OSError("media hash verification failed") + return data + + +def persist_inline_image_refs( + history: list[dict], + store: ImageMediaStore, +) -> list[dict]: + """Replace newly created data URIs with durable references before saving. + + Args: + history: Serialized messages about to be persisted. + store: Durable store for this application data root. + + Returns: + A new history list. Non-inline URLs remain unchanged for compatibility. + """ + import copy + + result = copy.deepcopy(history) + for message in result: + parts = message.get("content") if isinstance(message, dict) else None + if not isinstance(parts, list): + continue + for index, part in enumerate(parts): + if not isinstance(part, dict) or part.get("type") != "image_url": + continue + if part.get("_no_save"): + continue + image_url = part.get("image_url") + url = image_url.get("url") if isinstance(image_url, dict) else None + if not isinstance(url, str) or not url.startswith("data:image/"): + continue + header, payload = url.split(",", 1) + mime_type = header[5:].split(";", 1)[0] + data = base64.b64decode(payload, validate=True) + ref = store.put( + data, + mime_type, + image_url.get("detail") if isinstance(image_url, dict) else None, + image_url.get("id") if isinstance(image_url, dict) else None, + ) + parts[index] = ref.model_dump() + return result + + +async def materialize_image_media_refs( + contexts: list, + store: ImageMediaStore, + *, + strict: bool = False, +) -> list: + """Resolve only the selected references into a request-local view. + + Args: + contexts: Messages selected by the context manager. + store: Store for the application's configured data root. + strict: Fail on missing media when preparing a rollback export. + + Returns: + Request-local messages with images, or the original list if no refs exist. + + Raises: + ValueError: A reference has invalid metadata. + MemoryError: Image materialization exhausts process resources. + """ + import asyncio + + from astrbot import logger + from astrbot.core.agent.message import ContentPart, Message + + output = [] + for message in contexts: + serialized = message.model_dump() if isinstance(message, Message) else message + parts = serialized.get("content") + if not isinstance(parts, list) or not any( + isinstance(part, dict) and part.get("type") == "image_media_ref" + for part in parts + ): + output.append(message) + continue + resolved = [] + for part in parts: + if not isinstance(part, dict) or part.get("type") != "image_media_ref": + resolved.append(part) + continue + try: + ref = ImageMediaRef( + part["media_id"], + part["mime_type"], + part.get("width"), + part.get("height"), + part["byte_size"], + part.get("detail"), + part.get("version", 1), + part.get("image_id"), + ) + payload = await asyncio.to_thread(store.read, ref, {ref.media_id}) + image_url = { + "url": f"data:{ref.mime_type};base64,{base64.b64encode(payload).decode('ascii')}", + } + del payload + if ref.detail is not None: + image_url["detail"] = ref.detail + if ref.image_id is not None: + image_url["id"] = ref.image_id + resolved.append({"type": "image_url", "image_url": image_url}) + except (KeyError, TypeError, ValueError, OSError): + if strict: + raise + logger.warning("A selected conversation image is unavailable") + resolved.append({"type": "text", "text": "[Image unavailable]"}) + if isinstance(message, Message): + provider_message = message.model_copy() + provider_message.content = [ + ContentPart.model_validate(part) for part in resolved + ] + else: + provider_message = {**message, "content": resolved} + output.append(provider_message) + return output diff --git a/astrbot/core/utils/media_utils.py b/astrbot/core/utils/media_utils.py index b527c4b6c2..ce25658132 100644 --- a/astrbot/core/utils/media_utils.py +++ b/astrbot/core/utils/media_utils.py @@ -7,6 +7,7 @@ import asyncio import base64 import binascii +import ctypes import errno import hashlib import io @@ -27,7 +28,7 @@ from aiohttp import ClientError from PIL import Image as PILImage -from PIL import ImageOps +from PIL import ImageOps, UnidentifiedImageError from astrbot import logger from astrbot.core.utils.astrbot_path import get_astrbot_temp_path @@ -47,6 +48,461 @@ # 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 +IMAGE_COMPRESS_DEFAULT_MAX_ENCODED_BYTES = 4 * 1024 * 1024 + +_WEBP_PRESERVE = object() +_WEBP_DECODER = None +_WEBP_DECODER_UNAVAILABLE = False +_WEBP_ADVANCED_DECODER_AVAILABLE = False +_WEBP_DECODER_ABI_VERSION = 0x0210 + + +class _WebPRgbaBuffer(ctypes.Structure): + """ctypes view of libwebp's externally-owned packed pixel buffer.""" + + _fields_ = [ + ("rgba", ctypes.POINTER(ctypes.c_ubyte)), + ("stride", ctypes.c_int), + ("size", ctypes.c_size_t), + ] + + +class _WebPYuvaBuffer(ctypes.Structure): + """ctypes view of libwebp's YUVA buffer union arm.""" + + _fields_ = [ + ("y", ctypes.POINTER(ctypes.c_ubyte)), + ("u", ctypes.POINTER(ctypes.c_ubyte)), + ("v", ctypes.POINTER(ctypes.c_ubyte)), + ("a", ctypes.POINTER(ctypes.c_ubyte)), + ("y_stride", ctypes.c_int), + ("u_stride", ctypes.c_int), + ("v_stride", ctypes.c_int), + ("a_stride", ctypes.c_int), + ("y_size", ctypes.c_size_t), + ("u_size", ctypes.c_size_t), + ("v_size", ctypes.c_size_t), + ("a_size", ctypes.c_size_t), + ] + + +class _WebPBufferUnion(ctypes.Union): + """ctypes view of libwebp's decoded buffer union.""" + + _fields_ = [ + ("rgba", _WebPRgbaBuffer), + ("yuva", _WebPYuvaBuffer), + ] + + +class _WebPDecBuffer(ctypes.Structure): + """ctypes view of libwebp's decoder output descriptor.""" + + _fields_ = [ + ("colorspace", ctypes.c_int), + ("width", ctypes.c_int), + ("height", ctypes.c_int), + ("is_external_memory", ctypes.c_int), + ("u", _WebPBufferUnion), + ("pad", ctypes.c_uint32 * 4), + ("private_memory", ctypes.POINTER(ctypes.c_ubyte)), + ] + + +class _WebPBitstreamFeatures(ctypes.Structure): + """ctypes view of libwebp's bitstream feature structure.""" + + _fields_ = [ + ("width", ctypes.c_int), + ("height", ctypes.c_int), + ("has_alpha", ctypes.c_int), + ("has_animation", ctypes.c_int), + ("format", ctypes.c_int), + ("pad", ctypes.c_uint32 * 5), + ] + + +class _WebPDecoderOptions(ctypes.Structure): + """ctypes view of libwebp's decoder options structure.""" + + _fields_ = [ + (name, ctypes.c_int) + for name in ( + "bypass_filtering", + "no_fancy_upsampling", + "use_cropping", + "crop_left", + "crop_top", + "crop_width", + "crop_height", + "use_scaling", + "scaled_width", + "scaled_height", + "use_threads", + "dithering_strength", + "flip", + "alpha_dithering_strength", + ) + ] + [("pad", ctypes.c_uint32 * 5)] + + +class _WebPDecoderConfig(ctypes.Structure): + """ctypes view of libwebp's advanced decoder configuration.""" + + _fields_ = [ + ("input", _WebPBitstreamFeatures), + ("output", _WebPDecBuffer), + ("options", _WebPDecoderOptions), + ] + + +class ImagePayloadTooLargeError(ValueError): + """Raised when an encoded image cannot fit the preparation budget.""" + + +def _get_webp_decoder(): + """Load the decoder already shipped with Pillow when it exposes C APIs. + + Returns: + A configured ctypes library, or ``None`` when the Pillow build does + not expose the decoder symbols. + """ + global _WEBP_ADVANCED_DECODER_AVAILABLE + global _WEBP_DECODER, _WEBP_DECODER_UNAVAILABLE + if _WEBP_DECODER_UNAVAILABLE: + return None + if _WEBP_DECODER is not None: + return _WEBP_DECODER + try: + from PIL import _webp + + decoder = ctypes.CDLL(_webp.__file__) + get_info = decoder.WebPGetInfo + decode_rgb = decoder.WebPDecodeRGBInto + decode_rgba = decoder.WebPDecodeRGBAInto + get_info.argtypes = [ + ctypes.c_void_p, + ctypes.c_size_t, + ctypes.POINTER(ctypes.c_int), + ctypes.POINTER(ctypes.c_int), + ] + get_info.restype = ctypes.c_int + for decode in (decode_rgb, decode_rgba): + decode.argtypes = [ + ctypes.c_void_p, + ctypes.c_size_t, + ctypes.c_void_p, + ctypes.c_size_t, + ctypes.c_int, + ] + decode.restype = ctypes.c_void_p + except (AttributeError, ImportError, OSError): + _WEBP_DECODER_UNAVAILABLE = True + return None + try: + init_config = decoder.WebPInitDecoderConfigInternal + decode = decoder.WebPDecode + free_buffer = decoder.WebPFreeDecBuffer + init_config.argtypes = [ + ctypes.POINTER(_WebPDecoderConfig), + ctypes.c_int, + ] + init_config.restype = ctypes.c_int + decode.argtypes = [ + ctypes.c_void_p, + ctypes.c_size_t, + ctypes.POINTER(_WebPDecoderConfig), + ] + decode.restype = ctypes.c_int + free_buffer.argtypes = [ctypes.POINTER(_WebPDecBuffer)] + free_buffer.restype = None + except AttributeError: + _WEBP_ADVANCED_DECODER_AVAILABLE = False + else: + _WEBP_ADVANCED_DECODER_AVAILABLE = True + _WEBP_DECODER = decoder + return decoder + + +def _webp_container_properties( + data: bytes, +) -> tuple[int, int, bool, bool, bool] | None: + """Read WebP dimensions and flags without constructing a Pillow decoder. + + Args: + data: Complete WebP file bytes. + + Returns: + Width, height, alpha flag, animation flag, and EXIF flag. + ``None`` when the container needs the optional decoder to expose its + dimensions and that decoder is unavailable. + + Raises: + ValueError: The RIFF/WebP container is malformed. + """ + if len(data) < 20 or data[:4] != b"RIFF" or data[8:12] != b"WEBP": + raise ValueError("invalid WebP container") + width = height = None + has_alpha = False + has_animation = False + has_exif = False + offset = 12 + while offset + 8 <= len(data): + chunk_type = data[offset : offset + 4] + chunk_size = int.from_bytes(data[offset + 4 : offset + 8], "little") + start = offset + 8 + end = start + chunk_size + if end > len(data): + raise ValueError("truncated WebP chunk") + chunk = data[start:end] + if chunk_type == b"VP8X" and len(chunk) >= 10: + has_alpha = bool(chunk[0] & 0x10) + has_animation = bool(chunk[0] & 0x02) + width = 1 + int.from_bytes(chunk[4:7] + b"\0", "little") + height = 1 + int.from_bytes(chunk[7:10] + b"\0", "little") + elif chunk_type in {b"ANIM", b"ANMF"}: + has_animation = True + elif chunk_type == b"ALPH": + has_alpha = True + elif chunk_type == b"EXIF": + has_exif = True + elif chunk_type == b"VP8L" and len(chunk) >= 5 and chunk[0] == 0x2F: + has_alpha = has_alpha or bool( + int.from_bytes(chunk[1:5], "little") & (1 << 28) + ) + offset = end + (chunk_size & 1) + if width is None or height is None: + decoder = _get_webp_decoder() + if decoder is None: + return None + get_info = decoder.WebPGetInfo + width_value, height_value = ctypes.c_int(), ctypes.c_int() + if not get_info( + ctypes.c_char_p(data), + len(data), + ctypes.byref(width_value), + ctypes.byref(height_value), + ): + raise ValueError("invalid WebP bitstream") + width, height = width_value.value, height_value.value + return width, height, has_alpha, has_animation, has_exif + + +def _decode_webp_with_advanced_config( + decoder, + data: bytes, + width: int, + height: int, + target: tuple[int, int], + has_alpha: bool, +): + """Decode a scaled static WebP into one caller-owned pixel buffer. + + Args: + decoder: Configured libwebp shared library. + data: Complete encoded WebP bytes. + width: Original image width. + height: Original image height. + target: Requested output dimensions. + has_alpha: Whether the decoded image needs an alpha channel. + + Returns: + A Pillow image when the advanced ABI is available and decoding succeeds; + otherwise ``None`` so the caller can use the basic decoder path. + + Raises: + MemoryError: The caller-owned output buffer cannot be allocated. + """ + if not _WEBP_ADVANCED_DECODER_AVAILABLE or target == (width, height): + return None + + channels = 4 if has_alpha else 3 + mode = "RGBA" if has_alpha else "RGB" + pixels = bytearray(target[0] * target[1] * channels) + pixel_buffer = (ctypes.c_ubyte * len(pixels)).from_buffer(pixels) + config = _WebPDecoderConfig() + if not decoder.WebPInitDecoderConfigInternal( + ctypes.byref(config), _WEBP_DECODER_ABI_VERSION + ): + return None + config.output.colorspace = 1 if has_alpha else 0 + config.output.is_external_memory = 1 + config.output.u.rgba.rgba = ctypes.cast( + pixel_buffer, ctypes.POINTER(ctypes.c_ubyte) + ) + config.output.u.rgba.stride = target[0] * channels + config.output.u.rgba.size = len(pixels) + config.options.use_scaling = 1 + config.options.scaled_width, config.options.scaled_height = target + try: + if ( + decoder.WebPDecode(ctypes.c_char_p(data), len(data), ctypes.byref(config)) + != 0 + ): + return None + if (config.output.width, config.output.height) != target: + return None + finally: + decoder.WebPFreeDecBuffer(ctypes.byref(config.output)) + del pixel_buffer + image = PILImage.frombuffer(mode, target, pixels, "raw", mode, 0, 1) + image.format = "WEBP" + return image + + +def _open_static_webp( + source: bytes | Path, + max_size: int, + max_encoded_bytes: int, + *, + preserve_dimensions: bool = False, +): + """Decode static WebP directly into one caller-owned pixel buffer. + + Args: + source: Encoded WebP bytes or a local WebP path. + max_size: Maximum edge length for the eventual preparation step. + max_encoded_bytes: Base64 payload budget. + preserve_dimensions: Whether the eventual preparation step must retain + the original dimensions. + + Returns: + A Pillow image backed by one RGB/RGBA buffer, ``_WEBP_PRESERVE`` for a + compliant image, or ``None`` to use Pillow's compatibility path. + + Raises: + ImagePayloadTooLargeError: An animation already exceeds the budget. + ValueError: The WebP container is invalid. + MemoryError: The target pixel buffer cannot be allocated. + """ + data = source if isinstance(source, bytes) else source.read_bytes() + try: + properties = _webp_container_properties(data) + except ValueError: + return None + if properties is None: + return None + width, height, has_alpha, has_animation, has_exif = properties + PILImage._decompression_bomb_check((width, height)) + encoded_size = 4 * ((len(data) + 2) // 3) + if has_animation: + if encoded_size > max_encoded_bytes: + raise ImagePayloadTooLargeError( + f"Animated image exceeds the {max_encoded_bytes}-byte encoding limit" + ) + return _WEBP_PRESERVE + if encoded_size <= max_encoded_bytes and ( + preserve_dimensions or max(width, height) <= max_size + ): + return _WEBP_PRESERVE + # The direct route cannot carry EXIF orientation without loading Pillow's + # normal WebP plugin. Keep correctness for metadata-bearing images. + if has_exif: + return None + decoder = _get_webp_decoder() + if decoder is None: + return None + target = (width, height) + if not preserve_dimensions: + if max(width, height) > max_size: + scale = max_size / max(width, height) + target = ( + max(1, int(width * scale)), + max(1, int(height * scale)), + ) + elif encoded_size > max_encoded_bytes * 2 and max(width, height) >= max_size: + target = ( + max(1, width * 3 // 4), + max(1, height * 3 // 4), + ) + scaled = _decode_webp_with_advanced_config( + decoder, + data, + width, + height, + target, + has_alpha, + ) + if scaled is not None: + del data + return scaled + + channels = 4 if has_alpha else 3 + pixels = bytearray(width * height * channels) + pixel_buffer = (ctypes.c_ubyte * len(pixels)).from_buffer(pixels) + decode = decoder.WebPDecodeRGBAInto if has_alpha else decoder.WebPDecodeRGBInto + if not decode( + ctypes.c_char_p(data), + len(data), + pixel_buffer, + len(pixels), + width * channels, + ): + raise ValueError("invalid WebP bitstream") + del pixel_buffer, data + mode = "RGBA" if has_alpha else "RGB" + image = PILImage.frombuffer(mode, (width, height), pixels, "raw", mode, 0, 1) + image.format = "WEBP" + return image + + +@dataclass(slots=True) +class ImagePreparationOptions: + """Options for the single provider-facing image preparation boundary.""" + + enabled: bool = True + max_size: int = IMAGE_COMPRESS_DEFAULT_MAX_SIZE + quality: int = IMAGE_COMPRESS_DEFAULT_QUALITY + optimize: bool = IMAGE_COMPRESS_DEFAULT_OPTIMIZE + max_encoded_bytes: int | None = IMAGE_COMPRESS_DEFAULT_MAX_ENCODED_BYTES + preserve_dimensions: bool = False + + +def get_image_preparation_options( + provider_settings: dict | None, +) -> ImagePreparationOptions: + """Build image preparation options from provider settings. + + Args: + provider_settings: Provider-level image preparation configuration. + + Returns: + Validated image preparation options using the standard defaults. + """ + if not isinstance(provider_settings, dict): + return ImagePreparationOptions() + + enabled = provider_settings.get("image_compress_enabled", True) + if not isinstance(enabled, bool): + enabled = True + raw_options = provider_settings.get("image_compress_options", {}) + options = raw_options if isinstance(raw_options, dict) else {} + + max_size = options.get("max_size", IMAGE_COMPRESS_DEFAULT_MAX_SIZE) + if not isinstance(max_size, int) or isinstance(max_size, bool): + max_size = IMAGE_COMPRESS_DEFAULT_MAX_SIZE + + quality = options.get("quality", IMAGE_COMPRESS_DEFAULT_QUALITY) + if not isinstance(quality, int) or isinstance(quality, bool): + quality = IMAGE_COMPRESS_DEFAULT_QUALITY + + max_encoded_bytes = options.get( + "max_encoded_bytes", IMAGE_COMPRESS_DEFAULT_MAX_ENCODED_BYTES + ) + if not isinstance(max_encoded_bytes, int) or isinstance(max_encoded_bytes, bool): + max_encoded_bytes = IMAGE_COMPRESS_DEFAULT_MAX_ENCODED_BYTES + optimize = options.get("optimize", IMAGE_COMPRESS_DEFAULT_OPTIMIZE) + if not isinstance(optimize, bool): + optimize = IMAGE_COMPRESS_DEFAULT_OPTIMIZE + + return ImagePreparationOptions( + enabled=enabled, + max_size=max(max_size, 1), + quality=min(max(quality, 1), 100), + optimize=optimize, + max_encoded_bytes=max(max_encoded_bytes, 1), + ) + MEDIA_MIME_EXTENSIONS = { "audio/wav": ".wav", @@ -125,6 +581,22 @@ """ +@dataclass(frozen=True, slots=True) +class ImagePreparationInput: + """Describe an image entering the shared preparation boundary. + + Args: + value: Image path, URL, data URI, base64 reference, or raw bytes. + source_kind: Logical producer used for diagnostics and experiments. + cleanup_paths: Temporary paths owned by the caller and released after + preparation finishes. + """ + + value: MediaRefStr | bytes + source_kind: str = "unknown" + cleanup_paths: tuple[Path, ...] = () + + @dataclass(slots=True) class ResolvedMediaData: """Base64 media bytes plus the metadata needed by provider payloads. @@ -138,6 +610,7 @@ class ResolvedMediaData: base64_data: str mime_type: str format: str | None = None + byte_size: int | None = None def to_bytes(self) -> bytes: """Decode the base64 payload, accepting missing padding.""" @@ -180,7 +653,7 @@ def read_bytes(self) -> bytes: def to_base64(self) -> str: """Read the resolved local file and return raw base64 data.""" - return base64.b64encode(self.read_bytes()).decode("utf-8") + return _encode_file_to_base64(self.path) def to_data_url(self) -> str: """Read the resolved local file and return a data URL.""" @@ -327,6 +800,39 @@ def _decode_base64_payload( raise ValueError(error_message) from exc +def _encode_file_to_base64(path: Path) -> str: + """Encode a file without retaining a second full raw-image copy. + + Args: + path: Local file to read. + + Returns: + Base64 text without a data-URI prefix or line breaks. + + Raises: + OSError: The file cannot be opened or read. + """ + encoded = io.StringIO() + remainder = b"" + chunk = b"" + block = b"" + chunk_size = 1023 * 1024 + with path.open("rb") as source: + while True: + chunk = source.read(chunk_size) + if not chunk: + break + block = remainder + chunk if remainder else chunk + complete_size = len(block) - len(block) % 3 + if complete_size: + encoded.write(base64.b64encode(block[:complete_size]).decode("ascii")) + remainder = block[complete_size:] + del chunk, block + if remainder: + encoded.write(base64.b64encode(remainder).decode("ascii")) + return encoded.getvalue() + + def describe_media_ref(media_ref: object | None) -> str: """Return a log-safe description of a media reference. @@ -916,29 +1422,89 @@ async def open( async def resolve_image_ref_to_base64_data( - image_ref: MediaRefStr, + image_ref: MediaRefStr | bytes, *, strict: bool = False, default_mime_type: str | None = "image/jpeg", + options: ImagePreparationOptions | None = None, ) -> ResolvedMediaData | None: - """Resolve an image reference to losslessly encoded base64 data. + """Resolve and prepare an image reference for a provider request. + + ``strict=False`` returns ``None`` for invalid images so provider payload + assembly can skip bad image refs without failing the whole request. Resource + and size-limit errors still propagate because returning the original image + would allow an invalid or oversized payload to reach a provider. - Only materializes the source and detects its MIME type; no - provider-specific format conversion, frame extraction, or montage is - performed here. Platform senders and generic request assembly rely on - this to encode image bytes without transforming their content. + Args: + image_ref: Local path, URL, data URI, base64 reference, or image bytes. + strict: Raise ordinary resolution errors instead of returning ``None``. + default_mime_type: MIME fallback for otherwise unidentified images. + options: Dimension, encoding, and cleanup policy for image preparation. - ``strict=False`` returns ``None`` for invalid images so payload - assembly can skip bad image refs without failing the whole request. + Returns: + Prepared provider-ready image data, or ``None`` for a safely skippable + invalid image when ``strict`` is false. + + Raises: + ImagePayloadTooLargeError: The image cannot fit the configured budget. + MemoryError: Image preparation exhausts process resources. + OSError: The source cannot be read or the prepared output cannot be written. """ - return await MediaResolver( - image_ref, - media_type="image", - default_suffix=".bin", - ).to_base64_data( - strict=strict, - default_mime_type=default_mime_type, - ) + try: + return await prepare_image_source( + image_ref, + options=options, + default_mime_type=default_mime_type, + ) + except (ImagePayloadTooLargeError, MemoryError): + raise + except (OSError, ValueError): + is_legacy_base64 = isinstance(image_ref, str) and image_ref.startswith( + "base64://" + ) + if isinstance(image_ref, str) and not is_legacy_base64: + is_reference_scheme = image_ref.startswith( + ("http://", "https://", "data:") + ) or is_file_uri(image_ref) + try: + path_exists = Path(image_ref).exists() + except OSError: + path_exists = False + if not is_reference_scheme and not path_exists: + try: + _decode_base64_payload( + "".join(image_ref.split()), + error_message="invalid bare base64 media payload", + validate=True, + ) + except ValueError: + pass + else: + is_legacy_base64 = True + if is_legacy_base64: + # Preserve the historical fallback for opaque legacy base64 refs, + # while still enforcing the configured request budget. + legacy = await MediaResolver( + image_ref, + media_type="image", + default_suffix=".bin", + ).to_base64_data( + strict=strict, + default_mime_type=default_mime_type, + ) + if ( + legacy + and options + and options.max_encoded_bytes is not None + and len(legacy.base64_data) > options.max_encoded_bytes + ): + raise ImagePayloadTooLargeError( + f"Image exceeds the {options.max_encoded_bytes}-byte encoding limit" + ) + return legacy + if strict: + raise + return None def _image_convert_cache_dir() -> Path: @@ -1467,10 +2033,12 @@ async def resolve_audio_ref_to_base64_data( async def resolve_media_ref_to_base64_data( - media_ref: MediaRefStr, + media_ref: MediaRefStr | bytes, *, media_type: str, strict: bool = False, + image_options: ImagePreparationOptions | None = None, + default_mime_type: str | None = "image/jpeg", ) -> ResolvedMediaData | None: """Resolve a media reference to base64 data through one shared entrypoint. @@ -1479,7 +2047,12 @@ async def resolve_media_ref_to_base64_data( """ if media_type == "image": - return await resolve_image_ref_to_base64_data(media_ref, strict=strict) + return await resolve_image_ref_to_base64_data( + media_ref, + strict=strict, + default_mime_type=default_mime_type, + options=image_options, + ) if media_type == "audio": return await resolve_audio_ref_to_base64_data(media_ref) @@ -2023,134 +2596,444 @@ async def extract_video_cover( raise Exception("ffmpeg not found") +def _resize_alpha_in_strips( + source: PILImage.Image, size: tuple[int, int] +) -> PILImage.Image: + """Resize transparency with bounded premultiplication buffers. + + Pillow premultiplies a complete RGBA image before filtering. Processing + overlapping strips keeps those temporary buffers proportional to width, + while retaining the Lanczos support around each output strip. + + Args: + source: Caller-owned image with an alpha channel. + size: Output dimensions. + + Returns: + A caller-owned resized image. + + Raises: + OSError: Pixel decoding or resizing fails. + MemoryError: Output or temporary pixel allocation fails. + """ + output = PILImage.new(source.mode, size) + try: + ratio = source.height / size[1] + halo = math.ceil(3 * max(1, ratio)) + 1 + # Align strip boundaries with source rows whenever the rational scale + # permits it, avoiding floating-point phase changes at those boundaries. + alignment = size[1] // math.gcd(source.height, size[1]) + strip_height = max(alignment, 32 // alignment * alignment) + if strip_height > 64: + strip_height = 32 + for top in range(0, size[1], strip_height): + bottom = min(top + strip_height, size[1]) + start = max(0, math.floor(top * ratio) - halo) + end = min(source.height, math.ceil(bottom * ratio) + halo) + with source.crop((0, start, source.width, end)) as strip: + with strip.resize( + (size[0], bottom - top), + PILImage.Resampling.LANCZOS, + box=(0, top * ratio - start, source.width, bottom * ratio - start), + ) as resized: + output.paste(resized, (0, top)) + return output + except BaseException: + output.close() + raise + + def _compress_image_sync( source: bytes | Path, temp_dir: Path, max_size: int, quality: int, optimize: bool, + max_encoded_bytes: int = IMAGE_COMPRESS_DEFAULT_MAX_ENCODED_BYTES, + *, + preserve_dimensions: bool = False, ) -> str | None: - """Run image compression synchronously via ``asyncio.to_thread``. + """Prepare one image without allocating base64 for candidate measurements. Args: - source: Encoded image bytes or a local path to open inside the worker. - temp_dir: Directory where the compressed image should be written. - max_size: Longest edge of the compressed image in pixels. - quality: JPEG output quality in the range 1-100. - optimize: Whether Pillow should optimize the saved image. + source: Encoded bytes or a local file opened inside the worker. + temp_dir: Directory for the caller-owned prepared file. + max_size: Maximum edge length when resizing is allowed. + quality: Initial JPEG quality, between 1 and 100. + optimize: Whether to optimize the encoder output. + max_encoded_bytes: Maximum base64 payload size, excluding its URI header. + preserve_dimensions: Preserve screenshot coordinates, even over max_size. Returns: - The compressed image path, or ``None`` when the image should be kept as-is. + A caller-owned output path, or None to preserve the original bytes. + + Raises: + ImagePayloadTooLargeError: No candidate meets the encoded-byte budget. + ValueError: A size or quality option is invalid. + OSError: Reading, decoding or writing the image fails. """ + if max_size < 1 or max_encoded_bytes < 1 or not 1 <= quality <= 100: + raise ValueError("Image dimensions, byte budget and quality must be positive") + source_bytes = len(source) if isinstance(source, bytes) else source.stat().st_size + encoded_size = 4 * ((source_bytes + 2) // 3) + direct_webp = None + if isinstance(source, bytes): + is_webp = source[:4] == b"RIFF" and source[8:12] == b"WEBP" + else: + with source.open("rb") as source_stream: + header = source_stream.read(12) + is_webp = ( + len(header) >= 12 and header[:4] == b"RIFF" and header[8:12] == b"WEBP" + ) + if is_webp: + direct_webp = _open_static_webp( + source, + max_size, + max_encoded_bytes, + preserve_dimensions=preserve_dimensions, + ) + if direct_webp is _WEBP_PRESERVE: + return None fp = io.BytesIO(source) if isinstance(source, bytes) else source - with PILImage.open(fp) as opened_img: - converted_img: PILImage.Image | None = None + opened = direct_webp if direct_webp is not None else PILImage.open(fp) + with opened: + animated = getattr(opened, "n_frames", 1) > 1 + # This baseline preserves animation; do not flatten it to fit a budget. + if animated: + if encoded_size > max_encoded_bytes: + raise ImagePayloadTooLargeError( + f"Animated image exceeds the {max_encoded_bytes}-byte encoding limit" + ) + return None + if encoded_size <= max_encoded_bytes and ( + preserve_dimensions or max(opened.size) <= max_size + ): + return None + temp_dir.mkdir(parents=True, exist_ok=True) + path: Path | None = None + best_size: int | None = None + success = False + # Resize before EXIF handling loads the pixels. JPEG thumbnailing can + # then use decoder-level downsampling instead of a full-size RGB buffer. + # A square bound is unchanged by EXIF rotations and reflections. + if ( + opened.format == "JPEG" + and not preserve_dimensions + and max(opened.size) > max_size + ): + opened.thumbnail( + (max_size, max_size), PILImage.Resampling.LANCZOS, reducing_gap=1.0 + ) + # In-place orientation avoids holding a second full oriented image. + ImageOps.exif_transpose(opened, in_place=True) + has_alpha = opened.mode in {"RGBA", "LA"} or ( + opened.mode == "P" and "transparency" in opened.info + ) + target_mode = "RGBA" if has_alpha else "RGB" + converted = opened.convert(target_mode) if opened.mode != target_mode else None + working = converted if converted is not None else opened try: - if ( - getattr(opened_img, "is_animated", False) - or getattr(opened_img, "n_frames", 1) > 1 - ): - return None - - working_img = opened_img - image_has_alpha = opened_img.mode in {"RGBA", "LA"} or ( - opened_img.mode == "P" and "transparency" in opened_img.info + if not preserve_dimensions and max(working.size) > max_size: + working.thumbnail((max_size, max_size), PILImage.Resampling.LANCZOS) + qualities = ( + [quality] + if has_alpha + else sorted( + { + quality, + *[value for value in (85, 70, 55, 40) if value < quality], + }, + reverse=True, + ) ) - output_format = "PNG" if image_has_alpha else "JPEG" - output_suffix = ".png" if image_has_alpha else ".jpg" - - if image_has_alpha and opened_img.mode != "RGBA": - converted_img = opened_img.convert("RGBA") - working_img = converted_img - elif not image_has_alpha and opened_img.mode != "RGB": - converted_img = opened_img.convert("RGB") - working_img = converted_img - assert working_img is not None - - if max(working_img.size) > max_size: - working_img.thumbnail((max_size, max_size), PILImage.Resampling.LANCZOS) - - save_path = ( - temp_dir / f"compressed_{generate_timestamp_id()}{output_suffix}" + source_format = str(opened.format or "").upper() + # If the encoded source is already more than twice the budget, a + # same-size candidate cannot be a useful memory-saving first step. + # Resize once before encoding so the source pixels and encoder + # buffers are not resident together at the original dimensions. + skip_current_candidate = ( + not preserve_dimensions + and encoded_size > max_encoded_bytes * 2 + and max(working.size) >= max_size ) - save_kwargs: dict[str, int | bool] = {"optimize": optimize} - if output_format == "JPEG": - save_kwargs["quality"] = quality - working_img.save(save_path, output_format, **save_kwargs) - logger.debug(f"Image compressed successfully: {save_path}") - return str(save_path) + while True: + if not skip_current_candidate: + if has_alpha: + formats = [("PNG", ".png", None)] + else: + formats = ( + [("PNG", ".png", None)] if source_format == "PNG" else [] + ) + [ + ("JPEG", ".jpg", candidate_quality) + for candidate_quality in qualities + ] + for output_format, suffix, candidate_quality in formats: + candidate_path: Path | None = None + try: + with tempfile.NamedTemporaryFile( + dir=temp_dir, + prefix="compressed_", + suffix=suffix, + delete=False, + ) as output: + candidate_path = Path(output.name) + kwargs = {"optimize": optimize} + if candidate_quality is not None: + kwargs["quality"] = candidate_quality + working.save(candidate_path, output_format, **kwargs) + candidate_size = candidate_path.stat().st_size + encoded_size = 4 * ((candidate_size + 2) // 3) + if encoded_size <= max_encoded_bytes and ( + best_size is None or candidate_size < best_size + ): + if path is not None: + path.unlink(missing_ok=True) + path = candidate_path + best_size = candidate_size + candidate_path = None + finally: + if candidate_path is not None: + candidate_path.unlink(missing_ok=True) + # Compare a lossless PNG with the highest fitting JPEG + # quality; do not lower quality merely to minimize bytes. + if output_format == "JPEG" and path is not None: + break + skip_current_candidate = False + if path is not None: + success = True + return str(path) + if preserve_dimensions or working.size == (1, 1): + raise ImagePayloadTooLargeError( + f"Image cannot fit the {max_encoded_bytes}-byte encoding limit" + ) + # One current candidate, replaced on disk; never retain all encodings. + target = ( + max(1, working.width * 3 // 4), + max(1, working.height * 3 // 4), + ) + if has_alpha: + # Match thumbnail's aspect-ratio rounding. Rounding each + # edge independently can stretch narrow or odd-sized images. + width, height = target + aspect = working.width / working.height + if width / height >= aspect: + width = max( + min( + math.floor(height * aspect), + math.ceil(height * aspect), + key=lambda value: abs(aspect - value / height), + ), + 1, + ) + else: + height = max( + min( + math.floor(width / aspect), + math.ceil(width / aspect), + key=lambda value: ( + 0 if value == 0 else abs(aspect - width / value) + ), + ), + 1, + ) + target = (width, height) + resized = _resize_alpha_in_strips(working, target) + working.close() + working = converted = resized + else: + working.thumbnail(target, PILImage.Resampling.LANCZOS) finally: - if converted_img is not None: - converted_img.close() + if converted is not None: + converted.close() + if not success and path is not None: + path.unlink(missing_ok=True) async def compress_image( url_or_path: str, max_size: int = IMAGE_COMPRESS_DEFAULT_MAX_SIZE, quality: int = IMAGE_COMPRESS_DEFAULT_QUALITY, + max_encoded_bytes: int = IMAGE_COMPRESS_DEFAULT_MAX_ENCODED_BYTES, + *, + optimize: bool = IMAGE_COMPRESS_DEFAULT_OPTIMIZE, + preserve_dimensions: bool = False, ) -> str: - """Compress large user-uploaded images. + """Prepare a local image, preserving compliant bytes. Args: - url_or_path: Image path or URL. - max_size: Longest edge of the compressed image in pixels. - quality: JPEG output quality in the range 1-100. + url_or_path: Local path or inline image; remote URLs remain unresolved. + max_size: Maximum edge length when resizing is allowed. + quality: Initial JPEG quality. + max_encoded_bytes: Maximum base64 payload size. + preserve_dimensions: Preserve oriented screenshot dimensions. Returns: - The compressed image path. Returns the original path if compression - fails or the source does not need compression. - """ - max_size = max(int(max_size), 1) - quality = min(max(int(quality), 1), 100) - optimize = IMAGE_COMPRESS_DEFAULT_OPTIMIZE - min_file_size_bytes = int(IMAGE_COMPRESS_DEFAULT_MIN_FILE_SIZE_MB * 1024 * 1024) - image_source: bytes | Path | None = None + The original reference or a caller-owned prepared file path. - def _exceeds_max_size(source: bytes | Path) -> bool: - try: - fp = io.BytesIO(source) if isinstance(source, bytes) else source - with PILImage.open(fp) as opened_img: - return max(opened_img.size) > max_size - except Exception: # noqa: BLE001 - return False - - # Skip compression for remote images and return the original value. - if url_or_path.startswith("http"): + Raises: + ImagePayloadTooLargeError: The image cannot meet the byte limit. + OSError: Image decoding or filesystem access fails. + """ + if url_or_path.startswith(("http://", "https://")): return url_or_path - elif url_or_path.startswith("data:image"): - _header, encoded = url_or_path.split(",", 1) - image_source = _decode_base64_payload( - encoded, - error_message="invalid image data URI payload", + if url_or_path.startswith("data:image"): + _, encoded = url_or_path.split(",", 1) + image_source: bytes | Path = _decode_base64_payload( + encoded, error_message="invalid image data URI payload" ) - if len(image_source) < min_file_size_bytes and not _exceeds_max_size( - image_source - ): - return url_or_path else: - local_path = Path(url_or_path) - if not local_path.exists(): + image_source = Path(url_or_path) + if not image_source.exists(): return url_or_path - if local_path.stat().st_size < min_file_size_bytes and not _exceeds_max_size( - local_path - ): - return url_or_path - image_source = local_path - if image_source is None: - return url_or_path + worker = asyncio.create_task( + asyncio.to_thread( + _compress_image_sync, + image_source, + Path(get_astrbot_temp_path()), + max(int(max_size), 1), + min(max(int(quality), 1), 100), + optimize, + max(int(max_encoded_bytes), 1), + preserve_dimensions=preserve_dimensions, + ) + ) + try: + compressed_path = await asyncio.shield(worker) + except asyncio.CancelledError: + # Cancellation cannot stop Pillow in a thread. Retain ownership until + # the worker finishes, then release its output without deleting inputs. + def cleanup_finished(done: asyncio.Task) -> None: + try: + output = done.result() + if output is not None: + Path(output).unlink(missing_ok=True) + except Exception: + logger.warning("Cancelled image preparation cleanup failed") - temp_dir = Path(get_astrbot_temp_path()) - temp_dir.mkdir(parents=True, exist_ok=True) + worker.add_done_callback(cleanup_finished) + raise + return compressed_path or url_or_path - # Offload the blocking image processing task to a thread. - compressed_path = await asyncio.to_thread( - _compress_image_sync, - image_source, - temp_dir, - max_size, - quality, - optimize, + +async def prepare_image_source( + image_ref: MediaRefStr | bytes | ImagePreparationInput, + *, + options: ImagePreparationOptions | None = None, + default_mime_type: str | None = "image/jpeg", +) -> ResolvedMediaData: + """Resolve and prepare any image reference for a provider request. + + Args: + image_ref: Local path, HTTP(S) URL, data URI, base64 reference, bare + base64 payload, or an ``ImagePreparationInput`` descriptor. + options: Optional preparation limits. ``None`` uses the standard budget. + default_mime_type: Fallback MIME type for otherwise unidentified images. + + Returns: + Provider-ready base64 data and its detected MIME type. + + Raises: + ImagePayloadTooLargeError: The image cannot fit the configured budget. + OSError: The source cannot be read or decoded. + ValueError: The source is not a valid image. + """ + selected = options or ImagePreparationOptions() + preparation_input = ( + image_ref + if isinstance(image_ref, ImagePreparationInput) + else ImagePreparationInput(image_ref) ) - return compressed_path or url_or_path + source_ref = preparation_input.value + + async def _prepare() -> ResolvedMediaData: + owned_source: Path | None = None + try: + resolved_source = source_ref + if isinstance(source_ref, bytes): + owned_source = _temp_media_path("image", ".bin") + await asyncio.to_thread(owned_source.write_bytes, source_ref) + resolved_source = str(owned_source) + async with MediaResolver( + resolved_source, media_type="image", default_suffix=".bin" + ).as_path() as resolved: + if not selected.enabled: + source_size = resolved.path.stat().st_size + if ( + selected.max_encoded_bytes is not None + and 4 * ((source_size + 2) // 3) > selected.max_encoded_bytes + ): + raise ImagePayloadTooLargeError( + f"Image exceeds the {selected.max_encoded_bytes}-byte encoding limit" + ) + mime_type = await detect_image_mime_type_async( + resolved.path, default_mime_type=None + ) + if not mime_type: + raise ValueError( + f"Invalid image file: {describe_media_ref(resolved_source)}" + ) + image_size = source_size + return ResolvedMediaData( + base64_data=await asyncio.to_thread( + _encode_file_to_base64, resolved.path + ), + mime_type=mime_type, + byte_size=image_size, + ) + try: + prepared_path = await compress_image( + str(resolved.path), + max_size=selected.max_size, + quality=selected.quality, + max_encoded_bytes=( + selected.max_encoded_bytes + if selected.max_encoded_bytes is not None + else 2**63 - 1 + ), + optimize=selected.optimize, + preserve_dimensions=selected.preserve_dimensions, + ) + except UnidentifiedImageError as exc: + raise ValueError( + f"Invalid image file: {describe_media_ref(source_ref)}" + ) from exc + output_path = Path(prepared_path) + try: + mime_type = await detect_image_mime_type_async( + output_path, default_mime_type=None + ) + image_size = output_path.stat().st_size + encoded_data = await asyncio.to_thread( + _encode_file_to_base64, output_path + ) + finally: + if output_path != resolved.path: + output_path.unlink(missing_ok=True) + if not mime_type: + mime_type = resolved.mime_type or default_mime_type + if not mime_type: + raise ValueError( + f"Invalid image file: {describe_media_ref(resolved_source)}" + ) + return ResolvedMediaData( + base64_data=encoded_data, + mime_type=mime_type, + byte_size=image_size, + ) + finally: + if owned_source is not None: + owned_source.unlink(missing_ok=True) + for cleanup_path in preparation_input.cleanup_paths: + cleanup_path.unlink(missing_ok=True) + + worker = asyncio.create_task(_prepare()) + try: + return await asyncio.shield(worker) + except asyncio.CancelledError: + # The worker owns the resolver context until Pillow and file reads exit. + worker.add_done_callback( + lambda done: done.exception() if not done.cancelled() else None + ) + raise diff --git a/astrbot/dashboard/api/conversations.py b/astrbot/dashboard/api/conversations.py index f75003d457..06626524b7 100644 --- a/astrbot/dashboard/api/conversations.py +++ b/astrbot/dashboard/api/conversations.py @@ -3,7 +3,7 @@ from typing import Any, Literal from fastapi import APIRouter, Depends, Query, Request -from fastapi.responses import StreamingResponse +from fastapi.responses import Response, StreamingResponse from astrbot.dashboard.async_utils import run_maybe_async from astrbot.dashboard.responses import ApiError, ok @@ -173,6 +173,21 @@ async def export_conversations( return await _export_conversations(_model_dict(payload), service) +@router.get("/conversations/{conversation_id:path}/media/{media_id}") +async def preview_conversation_media( + conversation_id: str, + media_id: str, + user_id: str = Query(...), + _auth: AuthContext = Depends(require_data_scope), + service: ConversationService = Depends(get_service), +): + try: + media = await service.get_conversation_media(user_id, conversation_id, media_id) + return Response(content=media.data, media_type=media.mime_type) + except ConversationServiceError as exc: + _raise_conversation_error(exc) + + @router.post("/conversations/batch-delete") async def batch_delete_conversations( payload: ConversationBatchDeleteRequest, diff --git a/astrbot/dashboard/services/conversation_service.py b/astrbot/dashboard/services/conversation_service.py index 4018b2663f..4d6a2994a5 100644 --- a/astrbot/dashboard/services/conversation_service.py +++ b/astrbot/dashboard/services/conversation_service.py @@ -1,15 +1,23 @@ from __future__ import annotations +import asyncio import json import traceback from dataclasses import dataclass from datetime import datetime from io import BytesIO +from pathlib import Path from astrbot.core import logger from astrbot.core.core_lifecycle import AstrBotCoreLifecycle from astrbot.core.db import BaseDatabase from astrbot.core.umo_alias import build_umo_alias_map, parse_umo, serialize_umo_alias +from astrbot.core.utils.astrbot_path import get_astrbot_data_path +from astrbot.core.utils.image_media_store import ( + ImageMediaRef, + ImageMediaStore, + materialize_image_media_refs, +) class ConversationServiceError(Exception): @@ -23,6 +31,12 @@ class ConversationExport: mimetype: str = "application/jsonl" +@dataclass +class ConversationMedia: + data: bytes + mime_type: str + + class ConversationService: def __init__( self, @@ -32,6 +46,7 @@ def __init__( self.db_helper = db_helper self.conv_mgr = core_lifecycle.conversation_manager self.core_lifecycle = core_lifecycle + self.media_store = ImageMediaStore(Path(get_astrbot_data_path()) / "media") async def list_conversations( self, @@ -256,7 +271,9 @@ async def export_conversations(self, data: object) -> ConversationExport: continue webchat_titles = await self._get_webchat_titles([conversation]) - content = json.loads(conversation.history) + content = await materialize_image_media_refs( + json.loads(conversation.history), self.media_store, strict=True + ) export_record = { "cid": cid, "user_id": user_id, @@ -271,6 +288,8 @@ async def export_conversations(self, data: object) -> ConversationExport: } jsonl_lines.append(json.dumps(export_record, ensure_ascii=False)) exported_count += 1 + except MemoryError: + raise except Exception as exc: failed_items.append(f"user_id:{user_id}, cid:{cid} - {exc!s}") logger.error( @@ -290,6 +309,70 @@ async def export_conversations(self, data: object) -> ConversationExport: filename=f"astrbot_conversations_export_{timestamp}.jsonl", ) + async def get_conversation_media( + self, user_id: str, cid: str, media_id: str + ) -> ConversationMedia: + """Return one image referenced by an authorized conversation. + + Args: + user_id: Conversation owner and unified message origin. + cid: Conversation identifier. + media_id: Content hash from the persisted media reference. + + Returns: + Image bytes and their persisted MIME type. + + Raises: + ConversationServiceError: If ownership, reference, or media access fails. + """ + conversation = await self.db_helper.get_conversation_by_id(cid) + if not conversation: + raise ConversationServiceError("对话不存在") + if conversation.user_id != user_id: + raise ConversationServiceError("对话不存在") + try: + history = json.loads(conversation.history) + except (TypeError, json.JSONDecodeError) as exc: + raise ConversationServiceError("对话历史无效") from exc + refs = self._media_refs(history) + ref = refs.get(media_id) + if ref is None: + raise ConversationServiceError("媒体不存在") + try: + return ConversationMedia( + await asyncio.to_thread(self.media_store.read, ref, set(refs)), + ref.mime_type, + ) + except (OSError, PermissionError, FileNotFoundError) as exc: + raise ConversationServiceError("媒体暂时不可用,请从原始来源恢复") from exc + + @staticmethod + def _media_refs(value) -> dict[str, ImageMediaRef]: + found = {} + if isinstance(value, list): + for item in value: + found.update(ConversationService._media_refs(item)) + elif isinstance(value, dict): + if value.get("type") == "image_media_ref": + try: + ref = ImageMediaRef( + value["media_id"], + value["mime_type"], + value.get("width"), + value.get("height"), + value["byte_size"], + value.get("detail"), + value.get("version", 1), + value.get("image_id"), + ) + found[ref.media_id] = ref + except (KeyError, TypeError, ValueError): + pass + else: + for item in value.values(): + found.update(ConversationService._media_refs(item)) + return found + async def _delete_conversations(self, conversations: object) -> dict: if not isinstance(conversations, list) or not conversations: raise ConversationServiceError("批量删除时conversations参数不能为空") diff --git a/dashboard/src/api/generated/openapi-v1/sdk.gen.ts b/dashboard/src/api/generated/openapi-v1/sdk.gen.ts index 973b1971bb..7f9b52726f 100644 --- a/dashboard/src/api/generated/openapi-v1/sdk.gen.ts +++ b/dashboard/src/api/generated/openapi-v1/sdk.gen.ts @@ -1,7 +1,7 @@ // This file is auto-generated by @hey-api/openapi-ts import { createClient, createConfig, type OptionsLegacyParser, formDataBodySerializer } from '@hey-api/client-axios'; -import type { LoginData, LoginError, LoginResponse, LogoutError, LogoutResponse, GetAuthSetupStatusError, GetAuthSetupStatusResponse, SetupAuthData, SetupAuthError, SetupAuthResponse, SetupTotpData, SetupTotpError, SetupTotpResponse, RecoverTotpError, RecoverTotpResponse, UpdateAuthAccountData, UpdateAuthAccountError, UpdateAuthAccountResponse, ListApiKeysError, ListApiKeysResponse, CreateApiKeyData, CreateApiKeyError, CreateApiKeyResponse, RevokeApiKeyData, RevokeApiKeyError, RevokeApiKeyResponse, DeleteApiKeyData, DeleteApiKeyError, DeleteApiKeyResponse, GetSystemConfigSchemaError, GetSystemConfigSchemaResponse, GetSystemConfigError, GetSystemConfigResponse, UpdateSystemConfigData, UpdateSystemConfigError, UpdateSystemConfigResponse, GetSystemConfigRuntimeError, GetSystemConfigRuntimeResponse, GetConfigProfileSchemaError, GetConfigProfileSchemaResponse, ListConfigProfilesError, ListConfigProfilesResponse, CreateConfigProfileData, CreateConfigProfileError, CreateConfigProfileResponse, GetConfigProfileData, GetConfigProfileError, GetConfigProfileResponse, UpdateConfigProfileContentData, UpdateConfigProfileContentError, UpdateConfigProfileContentResponse, RenameConfigProfileData, RenameConfigProfileError, RenameConfigProfileResponse, DeleteConfigProfileData, DeleteConfigProfileError, DeleteConfigProfileResponse, ListConfigRoutesError, ListConfigRoutesResponse, ReplaceConfigRoutesData, ReplaceConfigRoutesError, ReplaceConfigRoutesResponse, UpsertConfigRouteData, UpsertConfigRouteError, UpsertConfigRouteResponse, DeleteConfigRouteData, DeleteConfigRouteError, DeleteConfigRouteResponse, ListBotTypesError, ListBotTypesResponse, RegisterBotTypeData, RegisterBotTypeError, RegisterBotTypeResponse, ListBotsData, ListBotsError, ListBotsResponse, CreateBotData, CreateBotError, CreateBotResponse, ListBotStatsError, ListBotStatsResponse, GetBotByIdData, GetBotByIdError, GetBotByIdResponse, UpdateBotByIdData, UpdateBotByIdError, UpdateBotByIdResponse, DeleteBotByIdData, DeleteBotByIdError, DeleteBotByIdResponse, SetBotEnabledByIdData, SetBotEnabledByIdError, SetBotEnabledByIdResponse, TestBotByIdData, TestBotByIdError, TestBotByIdResponse, GetBotData, GetBotError, GetBotResponse, UpdateBotData, UpdateBotError, UpdateBotResponse, DeleteBotData, DeleteBotError, DeleteBotResponse, SetBotEnabledData, SetBotEnabledError, SetBotEnabledResponse, TestBotData, TestBotError, TestBotResponse, GetProviderSchemaError, GetProviderSchemaResponse, ListProviderSourcesError, ListProviderSourcesResponse, CreateProviderSourceData, CreateProviderSourceError, CreateProviderSourceResponse, GetProviderSourceByIdData, GetProviderSourceByIdError, GetProviderSourceByIdResponse, UpsertProviderSourceByIdData, UpsertProviderSourceByIdError, UpsertProviderSourceByIdResponse, DeleteProviderSourceByIdData, DeleteProviderSourceByIdError, DeleteProviderSourceByIdResponse, ListProviderSourceModelsByIdData, ListProviderSourceModelsByIdError, ListProviderSourceModelsByIdResponse, ListProvidersBySourceIdData, ListProvidersBySourceIdError, ListProvidersBySourceIdResponse, CreateProviderInSourceByIdData, CreateProviderInSourceByIdError, CreateProviderInSourceByIdResponse, GetProviderSourceData, GetProviderSourceError, GetProviderSourceResponse, UpsertProviderSourceData, UpsertProviderSourceError, UpsertProviderSourceResponse, DeleteProviderSourceData, DeleteProviderSourceError, DeleteProviderSourceResponse, ListProviderSourceModelsData, ListProviderSourceModelsError, ListProviderSourceModelsResponse, ListProvidersBySourceData, ListProvidersBySourceError, ListProvidersBySourceResponse, CreateProviderInSourceData, CreateProviderInSourceError, CreateProviderInSourceResponse, ListProvidersData, ListProvidersError, ListProvidersResponse, CreateProviderData, CreateProviderError, CreateProviderResponse, GetProviderByIdData, GetProviderByIdError, GetProviderByIdResponse, UpdateProviderByIdData, UpdateProviderByIdError, UpdateProviderByIdResponse, DeleteProviderByIdData, DeleteProviderByIdError, DeleteProviderByIdResponse, SetProviderEnabledByIdData, SetProviderEnabledByIdError, SetProviderEnabledByIdResponse, TestProviderByIdData, TestProviderByIdError, TestProviderByIdResponse, GetProviderEmbeddingDimensionByIdData, GetProviderEmbeddingDimensionByIdError, GetProviderEmbeddingDimensionByIdResponse, GetProviderData, GetProviderError, GetProviderResponse, UpdateProviderData, UpdateProviderError, UpdateProviderResponse, DeleteProviderData, DeleteProviderError, DeleteProviderResponse, SetProviderEnabledData, SetProviderEnabledError, SetProviderEnabledResponse, TestProviderData, TestProviderError, TestProviderResponse, GetProviderEmbeddingDimensionData, GetProviderEmbeddingDimensionError, GetProviderEmbeddingDimensionResponse, SendChatMessageData, SendChatMessageError, SendChatMessageResponse, OpenChatWebSocketData, OpenLiveChatWebSocketData, OpenUnifiedChatWebSocketData, ListChatSessionsData, ListChatSessionsError, ListChatSessionsResponse, CreateChatSessionData, CreateChatSessionError, CreateChatSessionResponse, BatchDeleteChatSessionsData, BatchDeleteChatSessionsError, BatchDeleteChatSessionsResponse, GetChatSessionData, GetChatSessionError, GetChatSessionResponse, UpdateChatSessionData, UpdateChatSessionError, UpdateChatSessionResponse, DeleteChatSessionData, DeleteChatSessionError, DeleteChatSessionResponse, StopChatSessionData, StopChatSessionError, StopChatSessionResponse, ResumeChatRunData, ResumeChatRunError, ResumeChatRunResponse, UpdateChatMessageData, UpdateChatMessageError, UpdateChatMessageResponse, RegenerateChatMessageData, RegenerateChatMessageError, RegenerateChatMessageResponse, ListChatConfigsError, ListChatConfigsResponse, CreateChatThreadData, CreateChatThreadError, CreateChatThreadResponse, GetChatThreadData, GetChatThreadError, GetChatThreadResponse, DeleteChatThreadData, DeleteChatThreadError, DeleteChatThreadResponse, SendChatThreadMessageData, SendChatThreadMessageError, SendChatThreadMessageResponse, ListChatProjectsError, ListChatProjectsResponse, CreateChatProjectData, CreateChatProjectError, CreateChatProjectResponse, GetChatProjectData, GetChatProjectError, GetChatProjectResponse, UpdateChatProjectData, UpdateChatProjectError, UpdateChatProjectResponse, DeleteChatProjectData, DeleteChatProjectError, DeleteChatProjectResponse, ListChatProjectSessionsData, ListChatProjectSessionsError, ListChatProjectSessionsResponse, ListChatProjectWorkspaceFilesData, ListChatProjectWorkspaceFilesError, ListChatProjectWorkspaceFilesResponse, GetChatProjectWorkspaceFileData, GetChatProjectWorkspaceFileError, GetChatProjectWorkspaceFileResponse, DownloadChatProjectWorkspaceFileData, DownloadChatProjectWorkspaceFileError, DownloadChatProjectWorkspaceFileResponse, AddChatProjectSessionData, AddChatProjectSessionError, AddChatProjectSessionResponse, RemoveChatProjectSessionData, RemoveChatProjectSessionError, RemoveChatProjectSessionResponse, SendImMessageData, SendImMessageError, SendImMessageResponse, ListImBotsError, ListImBotsResponse, UploadFileData, UploadFileError, UploadFileResponse, UploadOpenApiFileData, UploadOpenApiFileError, UploadOpenApiFileResponse, DownloadOpenApiFileData, DownloadOpenApiFileError, DownloadOpenApiFileResponse, GetFileByNameData, GetFileByNameError, GetFileByNameResponse, GetTokenFileData, GetTokenFileError, GetTokenFileResponse, GetAttachmentData, GetAttachmentError, GetAttachmentResponse, DeleteAttachmentData, DeleteAttachmentError, DeleteAttachmentResponse, DownloadAttachmentData, DownloadAttachmentError, DownloadAttachmentResponse, ListPluginsData, ListPluginsError, ListPluginsResponse, GetPluginByIdData, GetPluginByIdError, GetPluginByIdResponse, UninstallPluginByIdData, UninstallPluginByIdError, UninstallPluginByIdResponse, GetPluginConfigByIdData, GetPluginConfigByIdError, GetPluginConfigByIdResponse, UpdatePluginConfigByIdData, UpdatePluginConfigByIdError, UpdatePluginConfigByIdResponse, GetPluginConfigSchemaByIdData, GetPluginConfigSchemaByIdError, GetPluginConfigSchemaByIdResponse, ListPluginConfigFilesByIdData, ListPluginConfigFilesByIdError, ListPluginConfigFilesByIdResponse, UploadPluginConfigFilesByIdData, UploadPluginConfigFilesByIdError, UploadPluginConfigFilesByIdResponse, DeletePluginConfigFileByIdData, DeletePluginConfigFileByIdError, DeletePluginConfigFileByIdResponse, GetPluginReadmeByIdData, GetPluginReadmeByIdError, GetPluginReadmeByIdResponse, GetPluginChangelogByIdData, GetPluginChangelogByIdError, GetPluginChangelogByIdResponse, ReloadPluginByIdData, ReloadPluginByIdError, ReloadPluginByIdResponse, SetPluginEnabledByIdData, SetPluginEnabledByIdError, SetPluginEnabledByIdResponse, ListPluginPagesByIdData, ListPluginPagesByIdError, ListPluginPagesByIdResponse, GetPluginPageByIdData, GetPluginPageByIdError, GetPluginPageByIdResponse, GetPluginPageAssetByIdData, GetPluginPageAssetByIdError, GetPluginPageAssetByIdResponse, GetPluginData, GetPluginError, GetPluginResponse, UninstallPluginData, UninstallPluginError, UninstallPluginResponse, GetPluginConfigData, GetPluginConfigError, GetPluginConfigResponse, UpdatePluginConfigData, UpdatePluginConfigError, UpdatePluginConfigResponse, UpdatePluginLogLevelData, UpdatePluginLogLevelError, UpdatePluginLogLevelResponse, GetPluginConfigSchemaData, GetPluginConfigSchemaError, GetPluginConfigSchemaResponse, ListPluginConfigFilesData, ListPluginConfigFilesError, ListPluginConfigFilesResponse, UploadPluginConfigFilesData, UploadPluginConfigFilesError, UploadPluginConfigFilesResponse, DeletePluginConfigFileData, DeletePluginConfigFileError, DeletePluginConfigFileResponse, GetPluginReadmeData, GetPluginReadmeError, GetPluginReadmeResponse, GetPluginChangelogData, GetPluginChangelogError, GetPluginChangelogResponse, ReloadPluginData, ReloadPluginError, ReloadPluginResponse, BindPluginSourceData, BindPluginSourceError, BindPluginSourceResponse, SetPluginEnabledData, SetPluginEnabledError, SetPluginEnabledResponse, UpdatePluginData, UpdatePluginError, UpdatePluginResponse, UpdatePluginsData, UpdatePluginsError, UpdatePluginsResponse, CheckPluginVersionSupportData, CheckPluginVersionSupportError, CheckPluginVersionSupportResponse, ValidatePluginRepoData, ValidatePluginRepoError, ValidatePluginRepoResponse, ListFailedPluginsError, ListFailedPluginsResponse, UninstallFailedPluginData, UninstallFailedPluginError, UninstallFailedPluginResponse, ReloadFailedPluginData, ReloadFailedPluginError, ReloadFailedPluginResponse, InstallPluginFromGithubData, InstallPluginFromGithubError, InstallPluginFromGithubResponse, InstallPluginFromGitData, InstallPluginFromGitError, InstallPluginFromGitResponse, InstallPluginFromUrlData, InstallPluginFromUrlError, InstallPluginFromUrlResponse, InstallPluginFromUploadData, InstallPluginFromUploadError, InstallPluginFromUploadResponse, ListPluginMarketData, ListPluginMarketError, ListPluginMarketResponse, ListPluginMarketCategoriesError, ListPluginMarketCategoriesResponse, ListPluginSourcesError, ListPluginSourcesResponse, CreatePluginSourceData, CreatePluginSourceError, CreatePluginSourceResponse, ReplacePluginSourcesData, ReplacePluginSourcesError, ReplacePluginSourcesResponse, DeletePluginSourceData, DeletePluginSourceError, DeletePluginSourceResponse, DeletePluginSourceByIdData, DeletePluginSourceByIdError, DeletePluginSourceByIdResponse, ListPluginPagesData, ListPluginPagesError, ListPluginPagesResponse, GetPluginPageData, GetPluginPageError, GetPluginPageResponse, GetPluginPageAssetData, GetPluginPageAssetError, GetPluginPageAssetResponse, GetPluginPageBridgeSdkError, GetPluginPageBridgeSdkResponse, GetPluginExtensionRouteData, GetPluginExtensionRouteError, GetPluginExtensionRouteResponse, PostPluginExtensionRouteData, PostPluginExtensionRouteError, PostPluginExtensionRouteResponse, PutPluginExtensionRouteData, PutPluginExtensionRouteError, PutPluginExtensionRouteResponse, PatchPluginExtensionRouteData, PatchPluginExtensionRouteError, PatchPluginExtensionRouteResponse, DeletePluginExtensionRouteData, DeletePluginExtensionRouteError, DeletePluginExtensionRouteResponse, ListCommandsData, ListCommandsError, ListCommandsResponse, UpdateCommandData, UpdateCommandError, UpdateCommandResponse, ListCommandConflictsError, ListCommandConflictsResponse, ListToolsData, ListToolsError, ListToolsResponse, SetToolEnabledData, SetToolEnabledError, SetToolEnabledResponse, SetToolPermissionData, SetToolPermissionError, SetToolPermissionResponse, ListMcpServersError, ListMcpServersResponse, CreateMcpServerData, CreateMcpServerError, CreateMcpServerResponse, UpdateMcpServerByNameData, UpdateMcpServerByNameError, UpdateMcpServerByNameResponse, DeleteMcpServerByNameData, DeleteMcpServerByNameError, DeleteMcpServerByNameResponse, SetMcpServerEnabledByNameData, SetMcpServerEnabledByNameError, SetMcpServerEnabledByNameResponse, TestMcpServerByNameData, TestMcpServerByNameError, TestMcpServerByNameResponse, UpdateMcpServerData, UpdateMcpServerError, UpdateMcpServerResponse, DeleteMcpServerData, DeleteMcpServerError, DeleteMcpServerResponse, SetMcpServerEnabledData, SetMcpServerEnabledError, SetMcpServerEnabledResponse, TestMcpServerData, TestMcpServerError, TestMcpServerResponse, SyncModelScopeMcpServersData, SyncModelScopeMcpServersError, SyncModelScopeMcpServersResponse, ListSkillsData, ListSkillsError, ListSkillsResponse, UploadSkillData, UploadSkillError, UploadSkillResponse, UploadSkillsBatchData, UploadSkillsBatchError, UploadSkillsBatchResponse, UpdateSkillByNameData, UpdateSkillByNameError, UpdateSkillByNameResponse, DeleteSkillByNameData, DeleteSkillByNameError, DeleteSkillByNameResponse, DownloadSkillByNameData, DownloadSkillByNameError, DownloadSkillByNameResponse, ListSkillFilesByNameData, ListSkillFilesByNameError, ListSkillFilesByNameResponse, GetSkillFileByNameData, GetSkillFileByNameError, GetSkillFileByNameResponse, UpdateSkillFileByNameData, UpdateSkillFileByNameError, UpdateSkillFileByNameResponse, UpdateSkillData, UpdateSkillError, UpdateSkillResponse, DeleteSkillData, DeleteSkillError, DeleteSkillResponse, DownloadSkillData, DownloadSkillError, DownloadSkillResponse, ListSkillFilesData, ListSkillFilesError, ListSkillFilesResponse, GetSkillFileData, GetSkillFileError, GetSkillFileResponse, UpdateSkillFileData, UpdateSkillFileError, UpdateSkillFileResponse, ListNeoSkillCandidatesData, ListNeoSkillCandidatesError, ListNeoSkillCandidatesResponse, ListNeoSkillReleasesData, ListNeoSkillReleasesError, ListNeoSkillReleasesResponse, GetNeoSkillPayloadData, GetNeoSkillPayloadError, GetNeoSkillPayloadResponse, EvaluateNeoSkillCandidateData, EvaluateNeoSkillCandidateError, EvaluateNeoSkillCandidateResponse, PromoteNeoSkillCandidateData, PromoteNeoSkillCandidateError, PromoteNeoSkillCandidateResponse, RollbackNeoSkillReleaseData, RollbackNeoSkillReleaseError, RollbackNeoSkillReleaseResponse, SyncNeoSkillReleaseData, SyncNeoSkillReleaseError, SyncNeoSkillReleaseResponse, DeleteNeoSkillCandidateData, DeleteNeoSkillCandidateError, DeleteNeoSkillCandidateResponse, DeleteNeoSkillReleaseData, DeleteNeoSkillReleaseError, DeleteNeoSkillReleaseResponse, ListKnowledgeBasesData, ListKnowledgeBasesError, ListKnowledgeBasesResponse, CreateKnowledgeBaseData, CreateKnowledgeBaseError, CreateKnowledgeBaseResponse, GetKnowledgeBaseData, GetKnowledgeBaseError, GetKnowledgeBaseResponse, UpdateKnowledgeBaseData, UpdateKnowledgeBaseError, UpdateKnowledgeBaseResponse, DeleteKnowledgeBaseData, DeleteKnowledgeBaseError, DeleteKnowledgeBaseResponse, GetKnowledgeBaseStatsData, GetKnowledgeBaseStatsError, GetKnowledgeBaseStatsResponse, ListKnowledgeDocumentsData, ListKnowledgeDocumentsError, ListKnowledgeDocumentsResponse, UploadKnowledgeDocumentData, UploadKnowledgeDocumentError, UploadKnowledgeDocumentResponse, ImportKnowledgeDocumentsData, ImportKnowledgeDocumentsError, ImportKnowledgeDocumentsResponse, ImportKnowledgeDocumentFromUrlData, ImportKnowledgeDocumentFromUrlError, ImportKnowledgeDocumentFromUrlResponse, GetKnowledgeDocumentData, GetKnowledgeDocumentError, GetKnowledgeDocumentResponse, DeleteKnowledgeDocumentData, DeleteKnowledgeDocumentError, DeleteKnowledgeDocumentResponse, ListKnowledgeChunksData, ListKnowledgeChunksError, ListKnowledgeChunksResponse, DeleteKnowledgeChunkData, DeleteKnowledgeChunkError, DeleteKnowledgeChunkResponse, RetrieveKnowledgeBaseData, RetrieveKnowledgeBaseError, RetrieveKnowledgeBaseResponse, GetKnowledgeTaskData, GetKnowledgeTaskError, GetKnowledgeTaskResponse, GetPersonaTreeError, GetPersonaTreeResponse, ListPersonasData, ListPersonasError, ListPersonasResponse, CreatePersonaData, CreatePersonaError, CreatePersonaResponse, GetPersonaByIdData, GetPersonaByIdError, GetPersonaByIdResponse, UpdatePersonaByIdData, UpdatePersonaByIdError, UpdatePersonaByIdResponse, DeletePersonaByIdData, DeletePersonaByIdError, DeletePersonaByIdResponse, GetPersonaData, GetPersonaError, GetPersonaResponse, UpdatePersonaData, UpdatePersonaError, UpdatePersonaResponse, DeletePersonaData, DeletePersonaError, DeletePersonaResponse, ListPersonaFoldersData, ListPersonaFoldersError, ListPersonaFoldersResponse, CreatePersonaFolderData, CreatePersonaFolderError, CreatePersonaFolderResponse, UpdatePersonaFolderData, UpdatePersonaFolderError, UpdatePersonaFolderResponse, DeletePersonaFolderData, DeletePersonaFolderError, DeletePersonaFolderResponse, MovePersonaItemData, MovePersonaItemError, MovePersonaItemResponse, ReorderPersonaItemsData, ReorderPersonaItemsError, ReorderPersonaItemsResponse, ListSessionsData, ListSessionsError, ListSessionsResponse, ListActiveUmosError, ListActiveUmosResponse, ListSessionRulesData, ListSessionRulesError, ListSessionRulesResponse, UpsertSessionRuleData, UpsertSessionRuleError, UpsertSessionRuleResponse, DeleteSessionRulesData, DeleteSessionRulesError, DeleteSessionRulesResponse, BatchUpdateSessionProviderData, BatchUpdateSessionProviderError, BatchUpdateSessionProviderResponse, BatchUpdateSessionServiceData, BatchUpdateSessionServiceError, BatchUpdateSessionServiceResponse, ListSessionGroupsError, ListSessionGroupsResponse, CreateSessionGroupData, CreateSessionGroupError, CreateSessionGroupResponse, UpdateSessionGroupData, UpdateSessionGroupError, UpdateSessionGroupResponse, DeleteSessionGroupData, DeleteSessionGroupError, DeleteSessionGroupResponse, ListConversationsData, ListConversationsError, ListConversationsResponse, GetConversationFilterOptionsError, GetConversationFilterOptionsResponse, BatchDeleteConversationsData, BatchDeleteConversationsError, BatchDeleteConversationsResponse, GetConversationData, GetConversationError, GetConversationResponse, UpdateConversationData, UpdateConversationError, UpdateConversationResponse, DeleteConversationData, DeleteConversationError, DeleteConversationResponse, ReplaceConversationMessagesData, ReplaceConversationMessagesError, ReplaceConversationMessagesResponse, ExportConversationsData, ExportConversationsError, ExportConversationsResponse, GetStatsData, GetStatsError, GetStatsResponse, GetProviderTokenStatsData, GetProviderTokenStatsError, GetProviderTokenStatsResponse, GetVersionError, GetVersionResponse, GetPublicVersionsError, GetPublicVersionsResponse, GetFirstNoticeData, GetFirstNoticeError, GetFirstNoticeResponse, TestGhproxyConnectionData, TestGhproxyConnectionError, TestGhproxyConnectionResponse, ListChangelogVersionsError, ListChangelogVersionsResponse, GetChangelogData, GetChangelogError, GetChangelogResponse, GetStartTimeError, GetStartTimeResponse, GetStorageStatusError, GetStorageStatusResponse, CleanupStorageData, CleanupStorageError, CleanupStorageResponse, RestartCoreError, RestartCoreResponse, ListBackupsData, ListBackupsError, ListBackupsResponse, CreateBackupData, CreateBackupError, CreateBackupResponse, UploadBackupData, UploadBackupError, UploadBackupResponse, InitBackupUploadData, InitBackupUploadError, InitBackupUploadResponse, UploadBackupChunkData, UploadBackupChunkError, UploadBackupChunkResponse, CompleteBackupUploadData, CompleteBackupUploadError, CompleteBackupUploadResponse, AbortBackupUploadData, AbortBackupUploadError, AbortBackupUploadResponse, GetBackupProgressData, GetBackupProgressError, GetBackupProgressResponse, DownloadBackupData, DownloadBackupError, DownloadBackupResponse, RenameBackupData, RenameBackupError, RenameBackupResponse, DeleteBackupData, DeleteBackupError, DeleteBackupResponse, CheckBackupData, CheckBackupError, CheckBackupResponse, ImportBackupData, ImportBackupError, ImportBackupResponse, CheckUpdateError, CheckUpdateResponse, ListReleasesData, ListReleasesError, ListReleasesResponse, UpdateCoreData, UpdateCoreError, UpdateCoreResponse, UpdateDashboardData, UpdateDashboardError, UpdateDashboardResponse, GetUpdateProgressData, GetUpdateProgressError, GetUpdateProgressResponse, InstallPipPackageData, InstallPipPackageError, InstallPipPackageResponse, ListCronJobsData, ListCronJobsError, ListCronJobsResponse, CreateCronJobData, CreateCronJobError, CreateCronJobResponse, UpdateCronJobData, UpdateCronJobError, UpdateCronJobResponse, DeleteCronJobData, DeleteCronJobError, DeleteCronJobResponse, RunCronJobData, RunCronJobError, RunCronJobResponse, StreamLiveLogsError, StreamLiveLogsResponse, GetLogHistoryError, GetLogHistoryResponse, GetTraceSettingsError, GetTraceSettingsResponse, UpdateTraceSettingsData, UpdateTraceSettingsError, UpdateTraceSettingsResponse, ListT2iTemplatesError, ListT2iTemplatesResponse, CreateT2iTemplateData, CreateT2iTemplateError, CreateT2iTemplateResponse, GetActiveT2iTemplateError, GetActiveT2iTemplateResponse, SetActiveT2iTemplateData, SetActiveT2iTemplateError, SetActiveT2iTemplateResponse, ResetDefaultT2iTemplateError, ResetDefaultT2iTemplateResponse, GetT2iTemplateData, GetT2iTemplateError, GetT2iTemplateResponse, UpdateT2iTemplateData, UpdateT2iTemplateError, UpdateT2iTemplateResponse, DeleteT2iTemplateData, DeleteT2iTemplateError, DeleteT2iTemplateResponse, GetSubagentConfigError, GetSubagentConfigResponse, UpdateSubagentConfigData, UpdateSubagentConfigError, UpdateSubagentConfigResponse, ListSubagentAvailableToolsError, ListSubagentAvailableToolsResponse, VerifyPlatformWebhookData, VerifyPlatformWebhookError, VerifyPlatformWebhookResponse, ReceivePlatformWebhookData, ReceivePlatformWebhookError, ReceivePlatformWebhookResponse } from './types.gen'; +import type { LoginData, LoginError, LoginResponse, LogoutError, LogoutResponse, GetAuthSetupStatusError, GetAuthSetupStatusResponse, SetupAuthData, SetupAuthError, SetupAuthResponse, SetupTotpData, SetupTotpError, SetupTotpResponse, RecoverTotpError, RecoverTotpResponse, UpdateAuthAccountData, UpdateAuthAccountError, UpdateAuthAccountResponse, ListApiKeysError, ListApiKeysResponse, CreateApiKeyData, CreateApiKeyError, CreateApiKeyResponse, RevokeApiKeyData, RevokeApiKeyError, RevokeApiKeyResponse, DeleteApiKeyData, DeleteApiKeyError, DeleteApiKeyResponse, GetSystemConfigSchemaError, GetSystemConfigSchemaResponse, GetSystemConfigError, GetSystemConfigResponse, UpdateSystemConfigData, UpdateSystemConfigError, UpdateSystemConfigResponse, GetSystemConfigRuntimeError, GetSystemConfigRuntimeResponse, GetConfigProfileSchemaError, GetConfigProfileSchemaResponse, ListConfigProfilesError, ListConfigProfilesResponse, CreateConfigProfileData, CreateConfigProfileError, CreateConfigProfileResponse, GetConfigProfileData, GetConfigProfileError, GetConfigProfileResponse, UpdateConfigProfileContentData, UpdateConfigProfileContentError, UpdateConfigProfileContentResponse, RenameConfigProfileData, RenameConfigProfileError, RenameConfigProfileResponse, DeleteConfigProfileData, DeleteConfigProfileError, DeleteConfigProfileResponse, ListConfigRoutesError, ListConfigRoutesResponse, ReplaceConfigRoutesData, ReplaceConfigRoutesError, ReplaceConfigRoutesResponse, UpsertConfigRouteData, UpsertConfigRouteError, UpsertConfigRouteResponse, DeleteConfigRouteData, DeleteConfigRouteError, DeleteConfigRouteResponse, ListBotTypesError, ListBotTypesResponse, RegisterBotTypeData, RegisterBotTypeError, RegisterBotTypeResponse, ListBotsData, ListBotsError, ListBotsResponse, CreateBotData, CreateBotError, CreateBotResponse, ListBotStatsError, ListBotStatsResponse, GetBotByIdData, GetBotByIdError, GetBotByIdResponse, UpdateBotByIdData, UpdateBotByIdError, UpdateBotByIdResponse, DeleteBotByIdData, DeleteBotByIdError, DeleteBotByIdResponse, SetBotEnabledByIdData, SetBotEnabledByIdError, SetBotEnabledByIdResponse, TestBotByIdData, TestBotByIdError, TestBotByIdResponse, GetBotData, GetBotError, GetBotResponse, UpdateBotData, UpdateBotError, UpdateBotResponse, DeleteBotData, DeleteBotError, DeleteBotResponse, SetBotEnabledData, SetBotEnabledError, SetBotEnabledResponse, TestBotData, TestBotError, TestBotResponse, GetProviderSchemaError, GetProviderSchemaResponse, ListProviderSourcesError, ListProviderSourcesResponse, CreateProviderSourceData, CreateProviderSourceError, CreateProviderSourceResponse, GetProviderSourceByIdData, GetProviderSourceByIdError, GetProviderSourceByIdResponse, UpsertProviderSourceByIdData, UpsertProviderSourceByIdError, UpsertProviderSourceByIdResponse, DeleteProviderSourceByIdData, DeleteProviderSourceByIdError, DeleteProviderSourceByIdResponse, ListProviderSourceModelsByIdData, ListProviderSourceModelsByIdError, ListProviderSourceModelsByIdResponse, ListProvidersBySourceIdData, ListProvidersBySourceIdError, ListProvidersBySourceIdResponse, CreateProviderInSourceByIdData, CreateProviderInSourceByIdError, CreateProviderInSourceByIdResponse, GetProviderSourceData, GetProviderSourceError, GetProviderSourceResponse, UpsertProviderSourceData, UpsertProviderSourceError, UpsertProviderSourceResponse, DeleteProviderSourceData, DeleteProviderSourceError, DeleteProviderSourceResponse, ListProviderSourceModelsData, ListProviderSourceModelsError, ListProviderSourceModelsResponse, ListProvidersBySourceData, ListProvidersBySourceError, ListProvidersBySourceResponse, CreateProviderInSourceData, CreateProviderInSourceError, CreateProviderInSourceResponse, ListProvidersData, ListProvidersError, ListProvidersResponse, CreateProviderData, CreateProviderError, CreateProviderResponse, GetProviderByIdData, GetProviderByIdError, GetProviderByIdResponse, UpdateProviderByIdData, UpdateProviderByIdError, UpdateProviderByIdResponse, DeleteProviderByIdData, DeleteProviderByIdError, DeleteProviderByIdResponse, SetProviderEnabledByIdData, SetProviderEnabledByIdError, SetProviderEnabledByIdResponse, TestProviderByIdData, TestProviderByIdError, TestProviderByIdResponse, GetProviderEmbeddingDimensionByIdData, GetProviderEmbeddingDimensionByIdError, GetProviderEmbeddingDimensionByIdResponse, GetProviderData, GetProviderError, GetProviderResponse, UpdateProviderData, UpdateProviderError, UpdateProviderResponse, DeleteProviderData, DeleteProviderError, DeleteProviderResponse, SetProviderEnabledData, SetProviderEnabledError, SetProviderEnabledResponse, TestProviderData, TestProviderError, TestProviderResponse, GetProviderEmbeddingDimensionData, GetProviderEmbeddingDimensionError, GetProviderEmbeddingDimensionResponse, SendChatMessageData, SendChatMessageError, SendChatMessageResponse, OpenChatWebSocketData, OpenLiveChatWebSocketData, OpenUnifiedChatWebSocketData, ListChatSessionsData, ListChatSessionsError, ListChatSessionsResponse, CreateChatSessionData, CreateChatSessionError, CreateChatSessionResponse, BatchDeleteChatSessionsData, BatchDeleteChatSessionsError, BatchDeleteChatSessionsResponse, GetChatSessionData, GetChatSessionError, GetChatSessionResponse, UpdateChatSessionData, UpdateChatSessionError, UpdateChatSessionResponse, DeleteChatSessionData, DeleteChatSessionError, DeleteChatSessionResponse, StopChatSessionData, StopChatSessionError, StopChatSessionResponse, ResumeChatRunData, ResumeChatRunError, ResumeChatRunResponse, UpdateChatMessageData, UpdateChatMessageError, UpdateChatMessageResponse, RegenerateChatMessageData, RegenerateChatMessageError, RegenerateChatMessageResponse, ListChatConfigsError, ListChatConfigsResponse, CreateChatThreadData, CreateChatThreadError, CreateChatThreadResponse, GetChatThreadData, GetChatThreadError, GetChatThreadResponse, DeleteChatThreadData, DeleteChatThreadError, DeleteChatThreadResponse, SendChatThreadMessageData, SendChatThreadMessageError, SendChatThreadMessageResponse, ListChatProjectsError, ListChatProjectsResponse, CreateChatProjectData, CreateChatProjectError, CreateChatProjectResponse, GetChatProjectData, GetChatProjectError, GetChatProjectResponse, UpdateChatProjectData, UpdateChatProjectError, UpdateChatProjectResponse, DeleteChatProjectData, DeleteChatProjectError, DeleteChatProjectResponse, ListChatProjectSessionsData, ListChatProjectSessionsError, ListChatProjectSessionsResponse, ListChatProjectWorkspaceFilesData, ListChatProjectWorkspaceFilesError, ListChatProjectWorkspaceFilesResponse, GetChatProjectWorkspaceFileData, GetChatProjectWorkspaceFileError, GetChatProjectWorkspaceFileResponse, DownloadChatProjectWorkspaceFileData, DownloadChatProjectWorkspaceFileError, DownloadChatProjectWorkspaceFileResponse, AddChatProjectSessionData, AddChatProjectSessionError, AddChatProjectSessionResponse, RemoveChatProjectSessionData, RemoveChatProjectSessionError, RemoveChatProjectSessionResponse, SendImMessageData, SendImMessageError, SendImMessageResponse, ListImBotsError, ListImBotsResponse, UploadFileData, UploadFileError, UploadFileResponse, UploadOpenApiFileData, UploadOpenApiFileError, UploadOpenApiFileResponse, DownloadOpenApiFileData, DownloadOpenApiFileError, DownloadOpenApiFileResponse, GetFileByNameData, GetFileByNameError, GetFileByNameResponse, GetTokenFileData, GetTokenFileError, GetTokenFileResponse, GetAttachmentData, GetAttachmentError, GetAttachmentResponse, DeleteAttachmentData, DeleteAttachmentError, DeleteAttachmentResponse, DownloadAttachmentData, DownloadAttachmentError, DownloadAttachmentResponse, ListPluginsData, ListPluginsError, ListPluginsResponse, GetPluginByIdData, GetPluginByIdError, GetPluginByIdResponse, UninstallPluginByIdData, UninstallPluginByIdError, UninstallPluginByIdResponse, GetPluginConfigByIdData, GetPluginConfigByIdError, GetPluginConfigByIdResponse, UpdatePluginConfigByIdData, UpdatePluginConfigByIdError, UpdatePluginConfigByIdResponse, GetPluginConfigSchemaByIdData, GetPluginConfigSchemaByIdError, GetPluginConfigSchemaByIdResponse, ListPluginConfigFilesByIdData, ListPluginConfigFilesByIdError, ListPluginConfigFilesByIdResponse, UploadPluginConfigFilesByIdData, UploadPluginConfigFilesByIdError, UploadPluginConfigFilesByIdResponse, DeletePluginConfigFileByIdData, DeletePluginConfigFileByIdError, DeletePluginConfigFileByIdResponse, GetPluginReadmeByIdData, GetPluginReadmeByIdError, GetPluginReadmeByIdResponse, GetPluginChangelogByIdData, GetPluginChangelogByIdError, GetPluginChangelogByIdResponse, ReloadPluginByIdData, ReloadPluginByIdError, ReloadPluginByIdResponse, SetPluginEnabledByIdData, SetPluginEnabledByIdError, SetPluginEnabledByIdResponse, ListPluginPagesByIdData, ListPluginPagesByIdError, ListPluginPagesByIdResponse, GetPluginPageByIdData, GetPluginPageByIdError, GetPluginPageByIdResponse, GetPluginPageAssetByIdData, GetPluginPageAssetByIdError, GetPluginPageAssetByIdResponse, GetPluginData, GetPluginError, GetPluginResponse, UninstallPluginData, UninstallPluginError, UninstallPluginResponse, GetPluginConfigData, GetPluginConfigError, GetPluginConfigResponse, UpdatePluginConfigData, UpdatePluginConfigError, UpdatePluginConfigResponse, UpdatePluginLogLevelData, UpdatePluginLogLevelError, UpdatePluginLogLevelResponse, GetPluginConfigSchemaData, GetPluginConfigSchemaError, GetPluginConfigSchemaResponse, ListPluginConfigFilesData, ListPluginConfigFilesError, ListPluginConfigFilesResponse, UploadPluginConfigFilesData, UploadPluginConfigFilesError, UploadPluginConfigFilesResponse, DeletePluginConfigFileData, DeletePluginConfigFileError, DeletePluginConfigFileResponse, GetPluginReadmeData, GetPluginReadmeError, GetPluginReadmeResponse, GetPluginChangelogData, GetPluginChangelogError, GetPluginChangelogResponse, ReloadPluginData, ReloadPluginError, ReloadPluginResponse, BindPluginSourceData, BindPluginSourceError, BindPluginSourceResponse, SetPluginEnabledData, SetPluginEnabledError, SetPluginEnabledResponse, UpdatePluginData, UpdatePluginError, UpdatePluginResponse, UpdatePluginsData, UpdatePluginsError, UpdatePluginsResponse, CheckPluginVersionSupportData, CheckPluginVersionSupportError, CheckPluginVersionSupportResponse, ValidatePluginRepoData, ValidatePluginRepoError, ValidatePluginRepoResponse, ListFailedPluginsError, ListFailedPluginsResponse, UninstallFailedPluginData, UninstallFailedPluginError, UninstallFailedPluginResponse, ReloadFailedPluginData, ReloadFailedPluginError, ReloadFailedPluginResponse, InstallPluginFromGithubData, InstallPluginFromGithubError, InstallPluginFromGithubResponse, InstallPluginFromGitData, InstallPluginFromGitError, InstallPluginFromGitResponse, InstallPluginFromUrlData, InstallPluginFromUrlError, InstallPluginFromUrlResponse, InstallPluginFromUploadData, InstallPluginFromUploadError, InstallPluginFromUploadResponse, ListPluginMarketData, ListPluginMarketError, ListPluginMarketResponse, ListPluginMarketCategoriesError, ListPluginMarketCategoriesResponse, ListPluginSourcesError, ListPluginSourcesResponse, CreatePluginSourceData, CreatePluginSourceError, CreatePluginSourceResponse, ReplacePluginSourcesData, ReplacePluginSourcesError, ReplacePluginSourcesResponse, DeletePluginSourceData, DeletePluginSourceError, DeletePluginSourceResponse, DeletePluginSourceByIdData, DeletePluginSourceByIdError, DeletePluginSourceByIdResponse, ListPluginPagesData, ListPluginPagesError, ListPluginPagesResponse, GetPluginPageData, GetPluginPageError, GetPluginPageResponse, GetPluginPageAssetData, GetPluginPageAssetError, GetPluginPageAssetResponse, GetPluginPageBridgeSdkError, GetPluginPageBridgeSdkResponse, GetPluginExtensionRouteData, GetPluginExtensionRouteError, GetPluginExtensionRouteResponse, PostPluginExtensionRouteData, PostPluginExtensionRouteError, PostPluginExtensionRouteResponse, PutPluginExtensionRouteData, PutPluginExtensionRouteError, PutPluginExtensionRouteResponse, PatchPluginExtensionRouteData, PatchPluginExtensionRouteError, PatchPluginExtensionRouteResponse, DeletePluginExtensionRouteData, DeletePluginExtensionRouteError, DeletePluginExtensionRouteResponse, ListCommandsData, ListCommandsError, ListCommandsResponse, UpdateCommandData, UpdateCommandError, UpdateCommandResponse, ListCommandConflictsError, ListCommandConflictsResponse, ListToolsData, ListToolsError, ListToolsResponse, SetToolEnabledData, SetToolEnabledError, SetToolEnabledResponse, SetToolPermissionData, SetToolPermissionError, SetToolPermissionResponse, ListMcpServersError, ListMcpServersResponse, CreateMcpServerData, CreateMcpServerError, CreateMcpServerResponse, UpdateMcpServerByNameData, UpdateMcpServerByNameError, UpdateMcpServerByNameResponse, DeleteMcpServerByNameData, DeleteMcpServerByNameError, DeleteMcpServerByNameResponse, SetMcpServerEnabledByNameData, SetMcpServerEnabledByNameError, SetMcpServerEnabledByNameResponse, TestMcpServerByNameData, TestMcpServerByNameError, TestMcpServerByNameResponse, UpdateMcpServerData, UpdateMcpServerError, UpdateMcpServerResponse, DeleteMcpServerData, DeleteMcpServerError, DeleteMcpServerResponse, SetMcpServerEnabledData, SetMcpServerEnabledError, SetMcpServerEnabledResponse, TestMcpServerData, TestMcpServerError, TestMcpServerResponse, SyncModelScopeMcpServersData, SyncModelScopeMcpServersError, SyncModelScopeMcpServersResponse, ListSkillsData, ListSkillsError, ListSkillsResponse, UploadSkillData, UploadSkillError, UploadSkillResponse, UploadSkillsBatchData, UploadSkillsBatchError, UploadSkillsBatchResponse, UpdateSkillByNameData, UpdateSkillByNameError, UpdateSkillByNameResponse, DeleteSkillByNameData, DeleteSkillByNameError, DeleteSkillByNameResponse, DownloadSkillByNameData, DownloadSkillByNameError, DownloadSkillByNameResponse, ListSkillFilesByNameData, ListSkillFilesByNameError, ListSkillFilesByNameResponse, GetSkillFileByNameData, GetSkillFileByNameError, GetSkillFileByNameResponse, UpdateSkillFileByNameData, UpdateSkillFileByNameError, UpdateSkillFileByNameResponse, UpdateSkillData, UpdateSkillError, UpdateSkillResponse, DeleteSkillData, DeleteSkillError, DeleteSkillResponse, DownloadSkillData, DownloadSkillError, DownloadSkillResponse, ListSkillFilesData, ListSkillFilesError, ListSkillFilesResponse, GetSkillFileData, GetSkillFileError, GetSkillFileResponse, UpdateSkillFileData, UpdateSkillFileError, UpdateSkillFileResponse, ListNeoSkillCandidatesData, ListNeoSkillCandidatesError, ListNeoSkillCandidatesResponse, ListNeoSkillReleasesData, ListNeoSkillReleasesError, ListNeoSkillReleasesResponse, GetNeoSkillPayloadData, GetNeoSkillPayloadError, GetNeoSkillPayloadResponse, EvaluateNeoSkillCandidateData, EvaluateNeoSkillCandidateError, EvaluateNeoSkillCandidateResponse, PromoteNeoSkillCandidateData, PromoteNeoSkillCandidateError, PromoteNeoSkillCandidateResponse, RollbackNeoSkillReleaseData, RollbackNeoSkillReleaseError, RollbackNeoSkillReleaseResponse, SyncNeoSkillReleaseData, SyncNeoSkillReleaseError, SyncNeoSkillReleaseResponse, DeleteNeoSkillCandidateData, DeleteNeoSkillCandidateError, DeleteNeoSkillCandidateResponse, DeleteNeoSkillReleaseData, DeleteNeoSkillReleaseError, DeleteNeoSkillReleaseResponse, ListKnowledgeBasesData, ListKnowledgeBasesError, ListKnowledgeBasesResponse, CreateKnowledgeBaseData, CreateKnowledgeBaseError, CreateKnowledgeBaseResponse, GetKnowledgeBaseData, GetKnowledgeBaseError, GetKnowledgeBaseResponse, UpdateKnowledgeBaseData, UpdateKnowledgeBaseError, UpdateKnowledgeBaseResponse, DeleteKnowledgeBaseData, DeleteKnowledgeBaseError, DeleteKnowledgeBaseResponse, GetKnowledgeBaseStatsData, GetKnowledgeBaseStatsError, GetKnowledgeBaseStatsResponse, ListKnowledgeDocumentsData, ListKnowledgeDocumentsError, ListKnowledgeDocumentsResponse, UploadKnowledgeDocumentData, UploadKnowledgeDocumentError, UploadKnowledgeDocumentResponse, ImportKnowledgeDocumentsData, ImportKnowledgeDocumentsError, ImportKnowledgeDocumentsResponse, ImportKnowledgeDocumentFromUrlData, ImportKnowledgeDocumentFromUrlError, ImportKnowledgeDocumentFromUrlResponse, GetKnowledgeDocumentData, GetKnowledgeDocumentError, GetKnowledgeDocumentResponse, DeleteKnowledgeDocumentData, DeleteKnowledgeDocumentError, DeleteKnowledgeDocumentResponse, ListKnowledgeChunksData, ListKnowledgeChunksError, ListKnowledgeChunksResponse, DeleteKnowledgeChunkData, DeleteKnowledgeChunkError, DeleteKnowledgeChunkResponse, RetrieveKnowledgeBaseData, RetrieveKnowledgeBaseError, RetrieveKnowledgeBaseResponse, GetKnowledgeTaskData, GetKnowledgeTaskError, GetKnowledgeTaskResponse, GetPersonaTreeError, GetPersonaTreeResponse, ListPersonasData, ListPersonasError, ListPersonasResponse, CreatePersonaData, CreatePersonaError, CreatePersonaResponse, GetPersonaByIdData, GetPersonaByIdError, GetPersonaByIdResponse, UpdatePersonaByIdData, UpdatePersonaByIdError, UpdatePersonaByIdResponse, DeletePersonaByIdData, DeletePersonaByIdError, DeletePersonaByIdResponse, GetPersonaData, GetPersonaError, GetPersonaResponse, UpdatePersonaData, UpdatePersonaError, UpdatePersonaResponse, DeletePersonaData, DeletePersonaError, DeletePersonaResponse, ListPersonaFoldersData, ListPersonaFoldersError, ListPersonaFoldersResponse, CreatePersonaFolderData, CreatePersonaFolderError, CreatePersonaFolderResponse, UpdatePersonaFolderData, UpdatePersonaFolderError, UpdatePersonaFolderResponse, DeletePersonaFolderData, DeletePersonaFolderError, DeletePersonaFolderResponse, MovePersonaItemData, MovePersonaItemError, MovePersonaItemResponse, ReorderPersonaItemsData, ReorderPersonaItemsError, ReorderPersonaItemsResponse, ListSessionsData, ListSessionsError, ListSessionsResponse, ListActiveUmosError, ListActiveUmosResponse, ListSessionRulesData, ListSessionRulesError, ListSessionRulesResponse, UpsertSessionRuleData, UpsertSessionRuleError, UpsertSessionRuleResponse, DeleteSessionRulesData, DeleteSessionRulesError, DeleteSessionRulesResponse, BatchUpdateSessionProviderData, BatchUpdateSessionProviderError, BatchUpdateSessionProviderResponse, BatchUpdateSessionServiceData, BatchUpdateSessionServiceError, BatchUpdateSessionServiceResponse, ListSessionGroupsError, ListSessionGroupsResponse, CreateSessionGroupData, CreateSessionGroupError, CreateSessionGroupResponse, UpdateSessionGroupData, UpdateSessionGroupError, UpdateSessionGroupResponse, DeleteSessionGroupData, DeleteSessionGroupError, DeleteSessionGroupResponse, ListConversationsData, ListConversationsError, ListConversationsResponse, GetConversationFilterOptionsError, GetConversationFilterOptionsResponse, BatchDeleteConversationsData, BatchDeleteConversationsError, BatchDeleteConversationsResponse, GetConversationData, GetConversationError, GetConversationResponse, UpdateConversationData, UpdateConversationError, UpdateConversationResponse, DeleteConversationData, DeleteConversationError, DeleteConversationResponse, ReplaceConversationMessagesData, ReplaceConversationMessagesError, ReplaceConversationMessagesResponse, PreviewConversationMediaData, PreviewConversationMediaError, PreviewConversationMediaResponse, ExportConversationsData, ExportConversationsError, ExportConversationsResponse, GetStatsData, GetStatsError, GetStatsResponse, GetProviderTokenStatsData, GetProviderTokenStatsError, GetProviderTokenStatsResponse, GetVersionError, GetVersionResponse, GetPublicVersionsError, GetPublicVersionsResponse, GetFirstNoticeData, GetFirstNoticeError, GetFirstNoticeResponse, TestGhproxyConnectionData, TestGhproxyConnectionError, TestGhproxyConnectionResponse, ListChangelogVersionsError, ListChangelogVersionsResponse, GetChangelogData, GetChangelogError, GetChangelogResponse, GetStartTimeError, GetStartTimeResponse, GetStorageStatusError, GetStorageStatusResponse, CleanupStorageData, CleanupStorageError, CleanupStorageResponse, RestartCoreError, RestartCoreResponse, ListBackupsData, ListBackupsError, ListBackupsResponse, CreateBackupData, CreateBackupError, CreateBackupResponse, UploadBackupData, UploadBackupError, UploadBackupResponse, InitBackupUploadData, InitBackupUploadError, InitBackupUploadResponse, UploadBackupChunkData, UploadBackupChunkError, UploadBackupChunkResponse, CompleteBackupUploadData, CompleteBackupUploadError, CompleteBackupUploadResponse, AbortBackupUploadData, AbortBackupUploadError, AbortBackupUploadResponse, GetBackupProgressData, GetBackupProgressError, GetBackupProgressResponse, DownloadBackupData, DownloadBackupError, DownloadBackupResponse, RenameBackupData, RenameBackupError, RenameBackupResponse, DeleteBackupData, DeleteBackupError, DeleteBackupResponse, CheckBackupData, CheckBackupError, CheckBackupResponse, ImportBackupData, ImportBackupError, ImportBackupResponse, CheckUpdateError, CheckUpdateResponse, ListReleasesData, ListReleasesError, ListReleasesResponse, UpdateCoreData, UpdateCoreError, UpdateCoreResponse, UpdateDashboardData, UpdateDashboardError, UpdateDashboardResponse, GetUpdateProgressData, GetUpdateProgressError, GetUpdateProgressResponse, InstallPipPackageData, InstallPipPackageError, InstallPipPackageResponse, ListCronJobsData, ListCronJobsError, ListCronJobsResponse, CreateCronJobData, CreateCronJobError, CreateCronJobResponse, UpdateCronJobData, UpdateCronJobError, UpdateCronJobResponse, DeleteCronJobData, DeleteCronJobError, DeleteCronJobResponse, RunCronJobData, RunCronJobError, RunCronJobResponse, StreamLiveLogsError, StreamLiveLogsResponse, GetLogHistoryError, GetLogHistoryResponse, GetTraceSettingsError, GetTraceSettingsResponse, UpdateTraceSettingsData, UpdateTraceSettingsError, UpdateTraceSettingsResponse, ListT2iTemplatesError, ListT2iTemplatesResponse, CreateT2iTemplateData, CreateT2iTemplateError, CreateT2iTemplateResponse, GetActiveT2iTemplateError, GetActiveT2iTemplateResponse, SetActiveT2iTemplateData, SetActiveT2iTemplateError, SetActiveT2iTemplateResponse, ResetDefaultT2iTemplateError, ResetDefaultT2iTemplateResponse, GetT2iTemplateData, GetT2iTemplateError, GetT2iTemplateResponse, UpdateT2iTemplateData, UpdateT2iTemplateError, UpdateT2iTemplateResponse, DeleteT2iTemplateData, DeleteT2iTemplateError, DeleteT2iTemplateResponse, GetSubagentConfigError, GetSubagentConfigResponse, UpdateSubagentConfigData, UpdateSubagentConfigError, UpdateSubagentConfigResponse, ListSubagentAvailableToolsError, ListSubagentAvailableToolsResponse, VerifyPlatformWebhookData, VerifyPlatformWebhookError, VerifyPlatformWebhookResponse, ReceivePlatformWebhookData, ReceivePlatformWebhookError, ReceivePlatformWebhookResponse } from './types.gen'; export const client = createClient(createConfig()); @@ -2616,6 +2616,16 @@ export const replaceConversationMessages = (options: OptionsLegacyParser) => { + return (options?.client ?? client).get({ + ...options, + url: '/api/v1/conversations/{conversation_id}/media/{media_id}' + }); +}; + /** * Export conversations */ diff --git a/dashboard/src/api/generated/openapi-v1/types.gen.ts b/dashboard/src/api/generated/openapi-v1/types.gen.ts index 37cad88028..8ee3f69094 100644 --- a/dashboard/src/api/generated/openapi-v1/types.gen.ts +++ b/dashboard/src/api/generated/openapi-v1/types.gen.ts @@ -3297,6 +3297,20 @@ export type ReplaceConversationMessagesResponse = (SuccessEnvelope); export type ReplaceConversationMessagesError = unknown; +export type PreviewConversationMediaData = { + path: { + conversation_id: string; + media_id: string; + }; + query: { + user_id: string; + }; +}; + +export type PreviewConversationMediaResponse = ((Blob | File)); + +export type PreviewConversationMediaError = (unknown); + export type ExportConversationsData = { body: ConversationExportRequest; }; diff --git a/dashboard/src/components/chat/AuthenticatedMediaImage.vue b/dashboard/src/components/chat/AuthenticatedMediaImage.vue new file mode 100644 index 0000000000..4b7e770aed --- /dev/null +++ b/dashboard/src/components/chat/AuthenticatedMediaImage.vue @@ -0,0 +1,70 @@ + + + diff --git a/dashboard/src/components/conversation/ConversationHistoryPreview.vue b/dashboard/src/components/conversation/ConversationHistoryPreview.vue index 11f660291d..86c8340ee2 100644 --- a/dashboard/src/components/conversation/ConversationHistoryPreview.vue +++ b/dashboard/src/components/conversation/ConversationHistoryPreview.vue @@ -4,8 +4,13 @@ import { ChevronRight, Minus, Plus } from "@lucide/vue"; import MarkdownIt from "markdown-it"; import DOMPurify from "dompurify"; import { useModuleI18n } from "@/i18n/composables"; +import AuthenticatedMediaImage from "@/components/chat/AuthenticatedMediaImage.vue"; -const props = defineProps<{ messages: unknown[] }>(); +const props = defineProps<{ + messages: unknown[]; + conversationId?: string; + userId?: string; +}>(); const { tm } = useModuleI18n("features/conversation"); const markdownEnabled = ref(true); const fontSize = ref(13); @@ -54,8 +59,18 @@ type PreviewPart = { text: string; html?: string; label?: string; + mediaId?: string; }; +function mediaUrl(mediaId: string) { + if (!props.conversationId || !props.userId) return ""; + return `/api/v1/conversations/${encodeURIComponent( + props.conversationId, + )}/media/${encodeURIComponent(mediaId)}?user_id=${encodeURIComponent( + props.userId, + )}`; +} + const records = computed(() => props.messages .filter( @@ -103,6 +118,16 @@ const records = computed(() => ) ) { parts.push({ kind: "image", text: item.image_url.url }); + } else if ( + item?.type === "image_media_ref" && + typeof item.media_id === "string" && + /^[0-9a-f]{64}$/i.test(item.media_id) + ) { + parts.push({ + kind: "image", + text: "", + mediaId: item.media_id, + }); } else { parts.push({ kind: "data", text: JSON.stringify(item, null, 2) }); } @@ -265,8 +290,18 @@ const records = computed(() =>
{{
               part.text
             }}
+ diff --git a/docs/en/dev/openapi-scopes.md b/docs/en/dev/openapi-scopes.md index 4bf87c6c4b..7b44c6e9b7 100644 --- a/docs/en/dev/openapi-scopes.md +++ b/docs/en/dev/openapi-scopes.md @@ -184,6 +184,7 @@ Manage conversations and platform-session data. | `GET` | `/api/v1/conversations/{conversation_id}` | — | | `PATCH` | `/api/v1/conversations/{conversation_id}` | — | | `DELETE` | `/api/v1/conversations/{conversation_id}` | — | +| `GET` | `/api/v1/conversations/{conversation_id}/media/{media_id}` | — | | `PUT` | `/api/v1/conversations/{conversation_id}/messages` | — | | `GET` | `/api/v1/session-groups` | — | | `POST` | `/api/v1/session-groups` | — | diff --git a/docs/public/openapi.json b/docs/public/openapi.json index b666e92afa..6a9d4a5f2c 100644 --- a/docs/public/openapi.json +++ b/docs/public/openapi.json @@ -6457,6 +6457,55 @@ "description": "**Required scope:** `data`" } }, + "/api/v1/conversations/{conversation_id}/media/{media_id}": { + "get": { + "tags": [ + "Conversations" + ], + "summary": "Preview an image referenced by a conversation", + "operationId": "previewConversationMedia", + "x-astrbot-scope": "data", + "parameters": [ + { + "$ref": "#/components/parameters/ConversationId" + }, + { + "name": "media_id", + "in": "path", + "required": true, + "schema": { + "type": "string", + "pattern": "^[0-9a-f]{64}$" + } + }, + { + "name": "user_id", + "in": "query", + "required": true, + "schema": { + "type": "string" + } + } + ], + "responses": { + "200": { + "description": "Image bytes", + "content": { + "image/*": { + "schema": { + "type": "string", + "format": "binary" + } + } + } + }, + "404": { + "description": "Conversation or media is unavailable" + } + }, + "description": "**Required scope:** `data`" + } + }, "/api/v1/conversations/export": { "post": { "tags": [ diff --git a/docs/zh/dev/openapi-scopes.md b/docs/zh/dev/openapi-scopes.md index 3152454eb3..36146c7ef8 100644 --- a/docs/zh/dev/openapi-scopes.md +++ b/docs/zh/dev/openapi-scopes.md @@ -184,6 +184,7 @@ outline: deep | `GET` | `/api/v1/conversations/{conversation_id}` | — | | `PATCH` | `/api/v1/conversations/{conversation_id}` | — | | `DELETE` | `/api/v1/conversations/{conversation_id}` | — | +| `GET` | `/api/v1/conversations/{conversation_id}/media/{media_id}` | — | | `PUT` | `/api/v1/conversations/{conversation_id}/messages` | — | | `GET` | `/api/v1/session-groups` | — | | `POST` | `/api/v1/session-groups` | — | diff --git a/openspec/changes/optimize-image-memory-lifecycle/.openspec.yaml b/openspec/changes/optimize-image-memory-lifecycle/.openspec.yaml new file mode 100644 index 0000000000..96db9a43b6 --- /dev/null +++ b/openspec/changes/optimize-image-memory-lifecycle/.openspec.yaml @@ -0,0 +1,2 @@ +schema: spec-driven +created: 2026-09-15 diff --git a/openspec/changes/optimize-image-memory-lifecycle/design.md b/openspec/changes/optimize-image-memory-lifecycle/design.md new file mode 100644 index 0000000000..47d11ea4a0 --- /dev/null +++ b/openspec/changes/optimize-image-memory-lifecycle/design.md @@ -0,0 +1,152 @@ +## Context + +See `proposal.md` and the three delta specs. The current provider request path turns image references into complete data URIs, while conversation saving serializes the model-visible message list. Image preparation currently focuses on longest-edge pixels and can pass through large pixel-compliant files. Existing histories already contain inline data URIs and must remain readable. + +## Goals / Non-Goals + +**Goals:** + +- Establish one preparation boundary shared by all image-producing inputs. +- Bound the encoded payload that can enter a provider request. +- Separate durable conversation records from media bytes and resolve media only for selected active messages. +- Keep old inline histories readable and make migration explicit. +- Measure file size, encoded payload size, request size, RSS high-water mark, Python allocations, temporary files, and post-request retention. + +**Non-Goals:** + +- Do not silently remove images from the active context merely to improve memory metrics. +- Do not rewrite existing conversations during ordinary reads. +- Do not assume the 5 MiB limit from #10089 applies to every provider. +- Do not make the provider-specific wire format depend on the persistence representation. + +## Decisions + +### 1. Use a media reference as the persistence boundary + +New history image parts use a versioned internal media reference containing a content hash, MIME type, dimensions, encoded byte size, and original image-detail metadata. Durable image objects live in a dedicated media directory under the configured AstrBot data root, NEVER the temporary directory or its age-based cleanup policy. Provider adapters continue to receive their existing image content shape after request-time materialization. No extra image-caption model is introduced. + +Alternative rejected: keep base64 in history and only trim it during compaction. This reduces some requests but leaves database growth, repeated JSON parsing, and compaction-record duplication. + +### 2. Keep backward-compatible dual readers + +The history reader accepts both the new reference form and existing inline `data:`/base64 forms. The writer emits references only for newly prepared images. Conversion of existing records is an explicit repair/migration operation with backup, dry-run statistics, hash verification, and rollback by restoring the original history. + +Alternative rejected: automatic in-place migration on first read. It makes a read destructive and risks losing recoverability when the media directory is unavailable. + +### 3. Make encoded bytes a first-class limit + +Image preparation uses independent limits for maximum dimensions and maximum encoded payload bytes. Compute base64 length as `4 * ceil(encoded_file_bytes / 3)` without allocating a base64 copy for every candidate. Retain only the current best candidate. Preserve already-compliant original bytes; otherwise try the original supported format and JPEG quality steps, select a valid candidate without enlarging compliant source content, and encode base64 only at the provider boundary. Preserve transparency when required. If no candidate fits, return a readable size error rather than silently returning the original oversized image. + +Proposed product defaults: preserve existing dimension and quality settings, add `image_compress_options.max_encoded_bytes = 4194304` (4 MiB of base64 per image). This is an AstrBot preparation budget, not a claim about vendor limits. A documented or configured smaller provider limit takes precedence. A known aggregate request limit is checked against the complete serialized request, independently of this per-image budget. Never reinterpret base64 bytes as text tokens. No universal aggregate vendor limit is invented. + +For CUA, preserve the oriented original pixel dimensions and coordinates; try encoding changes without resizing and return a readable error if the budget cannot be met. Do not use an enormous pixel limit as a substitute for an explicit preserve-dimensions flag. Animation follows the selected baseline's frame-selection behavior; an older branch that skips animations is not silently treated as the current montage implementation. Do not keep all decoded animation frames resident simultaneously. + +Alternative rejected: use only pixel dimensions or only source-file bytes. Neither predicts the final JSON payload reliably. + +### 4. Materialize only selected references + +Context selection happens before resolving media references. Images not selected for either the main request or a separate image-capable summary request are not opened, decoded, or encoded. The summary request and the subsequent main request have separate measured lifetimes. Selected references are materialized into a request-local structure and released after the request; the history model never receives the resulting data URI. In-flight cancellation must not delete a file still used by a non-cancelled image worker; cleanup runs when that worker actually exits. + +For a fixed provider configuration, prepared historical bytes are immutable. Do not re-encode the history when a new image arrives, tune old JPEG quality to the remaining request space, or substitute captions to achieve memory targets. Ordinary summary/truncation semantics are unchanged. A known oversized aggregate request fails clearly; this project does not add image eviction or a new compression trigger. B-only lazy loading MUST send byte-identical image content to the baseline. + +Alternative rejected: materialize all references first and let the context manager remove messages afterward. That preserves the current memory spike and defeats lazy loading. + +### 5. Treat all sources as adapters to one preparation function + +Platform attachments, quoted content, plugin/MCP results, file tools, and CUA screenshots normalize to a common input descriptor. The descriptor carries source ownership and cleanup responsibility. The preparation function owns format detection, dimension checks, byte checks, temporary-file cleanup, and diagnostic metadata. + +### 6. Use an external-process ablation harness + +Each measurement case runs in a fresh child process. The harness records RSS at high frequency, Python allocation snapshots separately, request JSON size, image byte sizes, media object counts, and temporary files. Cold-start, warm-start, single-image, multi-image, long-history, concurrent, restart, and missing-media cases are separate workloads. + +The implementation is accepted only when the optimized path preserves image count/order and provider-visible content for the active window. Memory reduction caused by silently dropping images is a failed experiment, not a success. + +## Risks / Trade-offs + +- [Risk] Media references can outlive their files. → Keep content hashes and metadata, report a bounded placeholder on misses, and add cleanup/reconciliation diagnostics. +- [Risk] Existing plugins may assume persisted `image_url` always contains a data URI. → Keep the dual reader and resolve references at the provider boundary; add compatibility tests for plugin/tool history. +- [Risk] Extra disk I/O can increase latency. → Resolve only selected images, cache prepared media by content hash within one request, and measure cold/warm latency separately. +- [Risk] Image conversion can use native memory invisible to `tracemalloc`. → Use an external RSS monitor and report both RSS and Python allocations. +- [Risk] A provider may have a smaller limit than the internal default. → Apply the effective provider limit during request preparation and preserve readable provider-specific errors. + +## Migration Plan + +1. Ship dual reading before new reference writing; factor switches belong in the experiment harness, not a new public rollout switch. +2. Run the ablation suite and compatibility tests before enabling reference writing by default. +3. Provide a dry-run migration report for old inline histories. +4. Migrate selected conversations only after media files and hashes are verified. +5. Before downgrading to a reader without reference support, stop writes and export/re-inline all reference-bearing histories into a verified backup, including conversations created after deployment. Merely disabling new writes does not make those records readable by old versions. Retain the media store until rollback verification completes. + +## 兼容与资源生命周期约束 + +- 当前基线工作区提交为 `1a2492a09f8722a9abd813f6a65fb687d53c88c9`,位于此前功能分支,不等同于 issue 的 v4.28.1。实现前固定目标基线及依赖,另建 v4.28.1 只读复现环境;不能拿两个版本的路径混合计算收益。 +- 新引用采用内部版本化图片部件,保留 image detail;对外 ProviderRequest 输入继续接受原来的路径、URL、data URI。新类型不得直接发送给远端 Provider。 +- Plugin/第三方 Provider 的兼容出口必须解析引用后再调用现有接口;内置 Provider 和摘要 Provider 都纳入测试。给插件的兼容视图可能仍需分配 base64,这部分成本如实报告,不能声称所有消费者都已懒加载。 +- 媒体文件采用内容哈希去重,先在同文件系统原子写入并核验,再提交消息引用。写入失败保持原会话不变;崩溃留下的未引用文件由维护命令处理。 +- 媒体读取通过内部 ID 查询并限制路径;WebUI 使用带会话权限校验的图片端点,不公开数据目录或仅凭哈希授权。会话列表/详情只返回引用和预览地址,不批量还原 base64。 +- 复用现有附件存储能力前核对其 WebChat 删除逻辑,禁止把 Agent 媒体直接挂入会被其他界面独立删除的生命周期。引用对象与消息内容绑定,跨会话去重不得带来跨会话读权限。 +- 首版不做运行中自动垃圾回收。显式维护命令在停止写入后扫描全部保留会话,输出 dry-run,隔离未引用对象;读写中断、解析失败时停止删除。移除或压缩一个会话不能删除仍被其他会话引用的文件。 +- 旧 inline 图片在普通读取时不转码、不改写;显式迁移只搬运原字节并核验哈希,不在迁移中夹带画质调整。历史包与媒体目录共同备份;导出保持自包含或附带媒体清单,导入必须校验引用。 +- 本地资源异常不应被笼统吞掉再走同一大请求;测试检查 MemoryError 保留异常类型并终止本次运行。413 与资源耗尽分开记录,不假设每个回退 Provider 的上限都相同,不引入自动多轮重压缩循环。 +- 这不是对话事件日志重构;当前正常上下文压缩及保存行为保持不变。媒体表示迁移与“保留所有原始对话”的产品需求不得混为一谈。 + +## 详细内存消融实验 + +### 因素与对照 + +三个开关仅供实验调用真实生产函数,不提供给用户。完整八组为 `000 / A00 / 0B0 / 00C / AB0 / A0C / 0BC / ABC`。 + +| 因素 | 唯一变化 | 必须保持不变 | +| --- | --- | --- | +| A 统一入口 | 所有来源进入同一准备链 | 使用原压缩算法和参数 | +| B 懒加载 | 同一准备结果外置存储,选窗后加载 | 图片字节、MIME、顺序、detail、消息位置 | +| C 压缩优化 | 仅在基线已有入口改变压缩算法 | 不顺便补齐工具/插件入口 | + +八组的零开关路径必须与固定基线的输入输出一致;用独立基线 checkout 的相同函数校验,不能把一个近似模拟器叫基线。另加纯文本对照,但不把它算作三因素收益。 + +### 固定样本 + +五类图片各包含普通与压力样本,每类三个固定种子,至少 30 个实例。文件在测量进程外生成,记录 SHA256、编码格式、尺寸、方向、透明度、帧数和字节数。 + +| 类别 | 普通 | 压力 | +| --- | --- | --- | +| PNG | 透明图/色块 | 1280×1280 高熵 RGBA,像素合规但字节超限 | +| JPEG | 常见照片、已压缩小文件 | 4000×3000 高细节、EXIF 旋转 | +| WebP | 静态有损 | 无损透明,动画分支另做覆盖 | +| GIF | 少帧动画 | 多帧动画,检查帧抽取/拼图峰值 | +| 工具截图 | 1920×1080 终端 | 4K 小字 UI 与坐标标记 | + +来源测试覆盖本地文件、HTTP、data URI、base64 URI,以及用户附件、引用、插件/MCP、FileRead、CUA。HTTP 使用本地服务器;这些入口调用现有适配代码,不能仅给同一函数贴不同来源名称。完整交叉用于单图工作负载,其余长会话用五种格式均衡混合。 + +### 工作负载 + +1. 单图单轮、同轮八张不同图片;分别统计每个阶段。 +2. 50 轮每轮一张不同图片,然后 20 轮纯文本;与重复同一张图片的独立测试区分去重收益。 +3. 10/50/100 张历史图;分别测全在窗口、按轮次裁剪、摘要成功和摘要失败四种状态。记录实际触发次数,不假设默认配置一定触发。 +4. 摘要模型分别支持和不支持图片。固定摘要返回值及用量,避免随机文本干扰;摘要图像输入遵循实际模态规则。 +5. 四会话并发、保存后重启、WebUI 只看详情/展开单张预览、显式导出;WebUI 与模型请求分别统计。 +6. 错误实验:缺失媒体、坏图片、取消、超时、写盘失败、5 MiB 聚合请求拒绝、独立单图限制、不同限制的回退 Provider。 +7. 静态图片循环固定窗口请求 200 次,记录 1/10/25/50/100/200 次,检查存活对象、临时文件、打开句柄及运行后残留;不强制 GC 的正常数据与诊断 GC 分开。 + +### 测量与防止实验污染 + +- 分阶段打点:读取数据库 → JSON 解析/消息构造 → 选窗 → 读取媒体 → 解码/缩放/编码 → base64 → Provider 组装 → HTTP 序列化/发送 → 保存 → 清理。另标记摘要请求阶段。 +- 记录原文件与输出文件字节、base64 长度、完整 HTTP JSON 字节、读取/解码/转码/base64 次数及字节量、阶段耗时、临时文件和持久化字节。磁盘数据不等价于内存收益。 +- 外部监测进程每 10 ms 采集被测进程及子进程 RSS;记录 OS 高水位。Windows 另采私有内存与工作集。两个平台单独比较,不直接比较绝对值。 +- 标准轮关闭 tracemalloc;诊断轮单独启用 tracemalloc,Linux 关键图另外追踪原生分配。先前约 24 MB WebP 结果缺乏初始化隔离,不能用作验收基线。 +- 每个实验单元十次独立进程运行,冷启动/预热三次后的热运行分别记录。按种子配对,随机交错八组,保存完整运行顺序。固定机器、解释器、依赖和处理并发数。 +- 本地模型服务返回固定文本、固定用量和固定摘要;真实客户端/SDK 必须参与最终请求组装。服务端流式计数,不保留所有历史请求体;结构校验另跑,不污染测量峰值。 +- 监测器不拷贝 base64、不输出原图、会话正文或密钥。样本生成、结果分析和监测服务不得计入被测进程。 +- 单次 120 秒或被测进程总内存达到 `min(2 GiB, 启动时可用内存的 25%)` 时停止并记录超限。崩溃、超时、OOM 数据不得丢弃或用零内存填充。 + +### 分析和验收 + +- 单因素效果比较 A/B/C 与基线,边际效果比较 ABC 与 BC/AC/AB;报告配对中位差、范围与原始样本,不把三个百分比相加。十次重复不宣称精确的 P99。 +- B-only 校验图片哈希、顺序、detail 和固定 Provider 的历史前缀;预期最终请求大小不变。请求时仍需装载整个有效图片窗口,必须披露这个下限。 +- C 校验方向、透明、截图文字和坐标;已合规原字节应原样保留,不能选用无必要膨胀的输出。超限失败不算“成功压缩”或内存收益。 +- CUA 的尺寸和坐标必须一致;失败不能伪装成空图成功。GIF 的既有抽帧语义必须一致。 +- B 在 50 张历史图的读取到选窗阶段,以相同消息窗口为前提,峰值增量降低至少 50%;ABC 在高熵长会话的成功可比场景中,整体峰值增量降低至少 30%。这是验收目标,不是已证明结果。 +- 小图正常场景峰值增量及中位耗时回退不超过 10%;图片质量或稳定历史前缀失败时,即使内存达标也不通过。 +- 持续请求后存活图片对象、文件句柄、临时文件不能按轮数线性增长。RSS 分配器保留不直接判定泄漏;结合对象与原生分配数据解释。 +- 实测平台最低覆盖 Linux/WSL 与原生 Windows,macOS 做功能回归。所有指标原始 JSONL/CSV、样本清单和曲线放在测试产物目录,不创建仓库 SUMMARY 报告。 +- 真实智谱只用于可选小样本画质/协议验证,与消融数据隔离。默认不发公网请求;密钥通过安全环境注入,绝不写入本规划或测试产物。 diff --git a/openspec/changes/optimize-image-memory-lifecycle/proposal.md b/openspec/changes/optimize-image-memory-lifecycle/proposal.md new file mode 100644 index 0000000000..8d6ffbcaa3 --- /dev/null +++ b/openspec/changes/optimize-image-memory-lifecycle/proposal.md @@ -0,0 +1,28 @@ +## Why + +AstrBot currently prepares images through several partially independent paths. Some inputs are resized only by pixel dimensions, while the final base64 payload can still exceed provider limits and consume several copies of the image in memory. Conversation history can also retain complete inline image data, so every later request pays the deserialization and request-construction cost again. Issue #10089 demonstrates request-size failures, while #10092 demonstrates process-level memory exhaustion. + +## What Changes + +- Add one image preparation contract for user attachments, quoted images, plugins, MCP/tool results, file reading, and screenshots. +- Enforce both dimension limits and a configurable final encoded-payload byte limit; reject or degrade images that cannot satisfy the limit instead of silently sending oversized data. +- Store historical images as durable media references and metadata rather than embedding complete base64 data in every conversation message. +- Resolve historical image references only when the active provider request needs them, with bounded loading and cleanup. +- Keep the current conversation semantics: images in the active context remain available to the model; this change does not silently remove images merely to improve memory numbers. +- Add memory-ablation tests covering preparation, history loading, request assembly, persistence, and repeated conversations. +- Preserve backward compatibility by reading existing inline data URIs and provide an explicit migration/repair path instead of rewriting them implicitly. + +## Capabilities + +### New Capabilities + +- `image-memory-lifecycle`: bounded image preparation, durable media references, lazy resolution, and memory-safe request assembly. +- `agent-context-image-budget`: active requests validate image byte size independently from token accounting and resolve references without changing image visibility. +- `conversation-history-media`: persisted messages may use media references and must remain readable across restart while preserving existing inline-image histories. + +## Impact + +- Affects media resolution and compression, provider request assembly, agent context processing, conversation persistence, tool/plugin image paths, and related tests. +- Adds a local media storage/index lifecycle with cleanup and missing-media handling. +- Adds configuration for final encoded image size and bounded media storage behavior; existing pixel/quality settings remain compatible. +- Provider payloads remain provider-specific; only the internal history representation and preparation boundary change. diff --git a/openspec/changes/optimize-image-memory-lifecycle/specs/agent-context-image-budget/spec.md b/openspec/changes/optimize-image-memory-lifecycle/specs/agent-context-image-budget/spec.md new file mode 100644 index 0000000000..0e14546564 --- /dev/null +++ b/openspec/changes/optimize-image-memory-lifecycle/specs/agent-context-image-budget/spec.md @@ -0,0 +1,40 @@ +## Purpose + +Keeps the active model context bounded by treating image payload bytes as a real request cost and by resolving only image data required by the active provider request. + +## ADDED Requirements + +### Requirement: Active context SHALL use a bounded image budget +The system SHALL validate encoded image bytes and known provider request-size limits independently from text token estimates. It MUST NOT introduce additional image eviction, replace available images with summaries, or re-encode historical images on every turn to fit a changing budget. + +#### Scenario: Context contains many images +- **WHEN** the active context contains images whose encoded payloads exceed the image budget +- **THEN** the system reports a readable size-limit error before sending a known-oversized request, without silently dropping images or changing historical image bytes + +### Requirement: Historical media SHALL be resolved on demand +The system SHALL keep historical image content as a resolvable media reference and SHALL materialize its bytes only when the image is selected for the active provider request. + +#### Scenario: Old image is outside the active context +- **WHEN** a historical image is outside the selected context window and is not input to the summarization request +- **THEN** the system does not read, decode, or base64-encode that image for the request + +### Requirement: Summarization SHALL retain its existing image input semantics +The system SHALL resolve selected image references for an image-capable summarization provider just as it does for the main provider. Images summarized out of the main window may still need loading for that separate summary request; this cost SHALL be measured separately. + +#### Scenario: Old image participates in a summary +- **WHEN** the existing context compression policy passes an old image to an image-capable summary provider +- **THEN** that provider receives the image rather than an unresolved reference, while the subsequent main request contains only its selected context + +### Requirement: Lazy loading SHALL preserve provider-visible image bytes +The system SHALL preserve image bytes, MIME type, detail, order, and message placement for a fixed prepared image across persistence, restart, and subsequent requests to the same provider configuration. + +#### Scenario: A later turn loads a persisted image +- **WHEN** a later text message causes an existing image reference to be materialized +- **THEN** the image content is byte-identical to its earlier prepared representation and the stable historical request prefix remains unchanged + +### Requirement: Existing inline images SHALL remain readable +The system SHALL continue to read existing persisted inline image data and SHALL apply the active image budget when such data is included in a request. + +#### Scenario: Existing conversation contains a data URI +- **WHEN** a conversation created before media references is loaded +- **THEN** the conversation remains readable and its inline image is handled by the same bounded request policy diff --git a/openspec/changes/optimize-image-memory-lifecycle/specs/conversation-history-media/spec.md b/openspec/changes/optimize-image-memory-lifecycle/specs/conversation-history-media/spec.md new file mode 100644 index 0000000000..23201fa078 --- /dev/null +++ b/openspec/changes/optimize-image-memory-lifecycle/specs/conversation-history-media/spec.md @@ -0,0 +1,37 @@ +## Purpose + +Separates durable conversation facts from large media payloads so that conversations remain restartable and image-capable without duplicating complete base64 data in every historical message. + +## ADDED Requirements + +### Requirement: New persisted image content SHALL use media references +The system SHALL persist a stable media reference with MIME type and image metadata for newly saved model-visible images instead of embedding the complete base64 payload in the conversation record. + +#### Scenario: Model-visible image is saved +- **WHEN** a request containing a newly prepared image is saved to conversation history +- **THEN** the history stores a media reference and sufficient metadata to resolve the image later + +### Requirement: Media references SHALL survive restart +The system SHALL resolve persisted media references after restart and SHALL report a readable missing-media result when the referenced media is unavailable. + +#### Scenario: Referenced media file is missing +- **WHEN** a conversation contains a reference whose media object no longer exists +- **THEN** loading the conversation does not crash, the UI reports unavailable media, and the model receives an explicit bounded missing-image placeholder; this is reported as a recovery condition, not a successful memory optimization + +### Requirement: Media lifetime SHALL be independent of temporary cleanup +Referenced media SHALL be durable outside the temporary-file cleanup domain. The system SHALL authorize image reads through conversation access, preserve shared media until no retained history references it, and never interpret a client-controlled media reference as an arbitrary local path. + +#### Scenario: Temporary files are cleaned after restart +- **WHEN** temporary media cleanup runs while a saved conversation still references an image +- **THEN** that saved image remains readable + +#### Scenario: Another conversation requests an unauthorized image +- **WHEN** a caller supplies an image identifier without access to its owning conversation +- **THEN** the request is rejected without exposing the image or its filesystem path + +### Requirement: History migration SHALL be explicit and reversible +The system SHALL preserve existing inline-image histories and SHALL provide an explicit migration or repair operation before replacing inline payloads with media references. + +#### Scenario: Existing history has inline base64 +- **WHEN** an old conversation is opened without migration +- **THEN** it remains readable and no destructive rewrite occurs automatically diff --git a/openspec/changes/optimize-image-memory-lifecycle/specs/image-memory-lifecycle/spec.md b/openspec/changes/optimize-image-memory-lifecycle/specs/image-memory-lifecycle/spec.md new file mode 100644 index 0000000000..871a21006e --- /dev/null +++ b/openspec/changes/optimize-image-memory-lifecycle/specs/image-memory-lifecycle/spec.md @@ -0,0 +1,26 @@ +## Purpose + +Provides one bounded lifecycle for images from users, tools, plugins, and external media so that model requests remain within provider limits without multiplying large image payloads in memory. + +## ADDED Requirements + +### Requirement: All model-bound images use one bounded preparation contract +The system SHALL apply the same preparation contract to user attachments, quoted images, plugin and MCP results, file-reading images, and computer-use screenshots before they become model image content. + +#### Scenario: Source-specific image enters a model request +- **WHEN** an image is supplied by any supported source +- **THEN** the system applies the same dimension, format, encoded-byte, cleanup, and failure rules before constructing provider content + +### Requirement: Image preparation SHALL enforce final encoded size +The system SHALL enforce a configurable upper bound on the final encoded image payload, in addition to pixel dimensions, and SHALL NOT silently send an image that exceeds the bound after preparation. + +#### Scenario: Pixel-compliant image exceeds the byte limit +- **WHEN** an image is within the configured dimensions but its encoded payload exceeds the configured byte limit +- **THEN** the system re-encodes or resizes it until it fits, or returns a readable image-too-large failure without sending the oversized payload + +### Requirement: Preparation SHALL release temporary resources +The system SHALL clean resolver-owned temporary files and release image-processing resources after request construction succeeds, fails, is cancelled, or times out. + +#### Scenario: Image preparation is cancelled +- **WHEN** request processing is cancelled during download, decode, resize, or encoding +- **THEN** temporary files and owned buffers are released without affecting unrelated media diff --git a/openspec/changes/optimize-image-memory-lifecycle/tasks.md b/openspec/changes/optimize-image-memory-lifecycle/tasks.md new file mode 100644 index 0000000000..22982703fa --- /dev/null +++ b/openspec/changes/optimize-image-memory-lifecycle/tasks.md @@ -0,0 +1,36 @@ +## 1. Baseline and measurement harness + +- [x] 1.1 Freeze the implementation and dependency baseline, create a separate read-only v4.28.1 reproduction environment for issues #10089 and #10092, and verify both environments report their commit and dependency versions. +- [x] 1.2 Add an external-process image benchmark harness that records stage timings, RSS high-water marks, Python allocation diagnostics, child-process memory, request JSON bytes, media bytes, temporary files, and failures without printing image data or secrets. +- [x] 1.3 Generate deterministic PNG, JPEG, WebP, GIF, and screenshot fixtures with ordinary and high-entropy cases, record hashes and metadata, and verify the fixture generator is outside the measured child process. +- [ ] 1.4 Implement the eight-factor ablation runner for unified-source preparation, lazy history loading, and compression optimization; verify paired seeds, cold/warm runs, ten repetitions, and raw JSONL/CSV output. + +## 2. Unified preparation and encoded-byte limits + +- [x] 2.1 Define the internal image preparation descriptor and one preparation entry point for platform attachments, quoted images, plugins/MCP, file tools, and CUA screenshots; verify every listed source reaches it with source ownership and cleanup metadata. +- [x] 2.2 Extend image preparation with an encoded payload byte limit independent from pixel dimensions; verify a pixel-compliant high-entropy PNG is re-encoded or rejected instead of passed through. +- [x] 2.3 Implement candidate-size measurement without retaining all candidate bytes and preserve transparency, EXIF orientation, CUA dimensions, animation behavior, and already-compliant bytes; verify format-specific fixtures and coordinate/visual checks. +- [x] 2.4 Add typed, readable image-size and resource errors and ensure oversize input is not silently returned to the Provider or retried through unrelated fallback providers; verify #10089's request-size scenario. +- [x] 2.5 Add cancellation, timeout, decode failure, and write failure cleanup tests; verify temporary files, image-library resources, and worker tasks are released on every path. + +## 3. Durable media references and lazy materialization + +- [x] 3.1 Add versioned durable media objects under the configured data directory with content-hash deduplication, atomic write/verify, MIME/dimension/byte/detail metadata, and conversation-scoped access checks; verify partial writes never create usable references. +- [x] 3.2 Add a persisted image-reference representation and dual history reader for new references plus existing inline data URIs/base64; verify old histories remain readable and new histories do not embed complete base64. +- [x] 3.3 Resolve selected references only after context selection and materialize them into a request-local provider view; verify out-of-window images are not opened/decoded/encoded and B-only provider-visible bytes, order, detail, and placement match baseline. +- [x] 3.4 Keep durable media outside temporary cleanup and add missing-media, restart, unauthorized-reference, session-delete, and cross-session shared-media tests; verify missing media produces a bounded recovery result without exposing paths. +- [x] 3.5 Add explicit dry-run migration/export/repair behavior for old inline histories with backup and hash verification; verify ordinary reads never rewrite history and rollback remains possible. + +## 4. Context and persistence integration + +- [x] 4.1 Integrate image byte validation with the active request and summarization-provider paths without adding image eviction or changing existing truncation/summary semantics; verify image-capable summaries receive selected images and the main request receives its selected context. +- [x] 4.2 Ensure history saving persists references rather than request-time data URIs while provider adapters continue receiving their existing wire shapes; verify plugin, tool, built-in Provider, WebUI detail, and export compatibility. +- [x] 4.3 Separate 413 request-size, image-too-large, missing-media, and MemoryError handling; verify MemoryError preserves its type and does not trigger repeated oversized fallback requests. + +## 5. Verification and rollout + +- [ ] 5.1 Run the full matrix across single image, eight images, 50-round history, 200 fixed-window requests, four-session concurrency, restart, WebUI preview, and error workloads; verify all metrics and raw failures are retained. +- [ ] 5.2 Compare A/B/C main and marginal effects with image-count/order/content invariants; verify memory gains are not caused by silently omitting active images. +- [ ] 5.3 Validate Linux/WSL and native Windows memory behavior, run macOS functional regression, and verify normal small-image latency/peak-memory regression stays within the documented target. +- [x] 5.4 Run `ruff format --check .`, `ruff check .`, focused image/context/history tests, and the relevant full test suites; verify no source comments/logs violate repository language rules. +- [ ] 5.5 Review migration backup/rollback and media cleanup behavior, document measured results in test artifacts outside repository SUMMARY files, and only then enable new reference writes by default. diff --git a/openspec/openapi-v1.yaml b/openspec/openapi-v1.yaml index ab63053bf0..fb9f3bba6f 100644 --- a/openspec/openapi-v1.yaml +++ b/openspec/openapi-v1.yaml @@ -4351,6 +4351,36 @@ paths: "200": $ref: "#/components/responses/Ok" + /api/v1/conversations/{conversation_id}/media/{media_id}: + get: + tags: [Conversations] + summary: Preview an image referenced by a conversation + operationId: previewConversationMedia + x-astrbot-scope: data + parameters: + - $ref: "#/components/parameters/ConversationId" + - name: media_id + in: path + required: true + schema: + type: string + pattern: "^[0-9a-f]{64}$" + - name: user_id + in: query + required: true + schema: + type: string + responses: + "200": + description: Image bytes + content: + image/*: + schema: + type: string + format: binary + "404": + description: Conversation or media is unavailable + /api/v1/conversations/export: post: tags: [Conversations] diff --git a/scripts/image_memory_bench/ablation_runner.py b/scripts/image_memory_bench/ablation_runner.py new file mode 100644 index 0000000000..9c8b08abb9 --- /dev/null +++ b/scripts/image_memory_bench/ablation_runner.py @@ -0,0 +1,66 @@ +"""Run the real 2^3 image experiment only when production factor flips exist. + +The current checkout does not expose independent A, B, and C experiment +switches. This runner therefore refuses to fabricate a matrix and emits an +auditable blocked manifest until each factor is supplied by a real checkout. +""" + +from __future__ import annotations + +import argparse +import itertools +import json +import subprocess +from pathlib import Path + + +def git_identity(repo: Path) -> dict[str, str]: + """Return commit and dirty diff identity for a checkout.""" + commit = subprocess.check_output( + ["git", "rev-parse", "HEAD"], cwd=repo, text=True + ).strip() + diff = subprocess.check_output(["git", "diff", "--binary"], cwd=repo) + import hashlib + + return { + "commit": commit, + "working_tree_diff_sha256": hashlib.sha256(diff).hexdigest(), + } + + +def main() -> int: + """Validate factor provenance and write a blocked or runnable plan.""" + parser = argparse.ArgumentParser() + parser.add_argument("output", type=Path) + parser.add_argument("--baseline-repo", type=Path, required=True) + parser.add_argument("--candidate-repo", type=Path, required=True) + parser.add_argument("--factor-entrypoint", type=Path) + args = parser.parse_args() + factors = { + "A": "unavailable: no independent production source-preparation entrypoint supplied", + "B": "unavailable: no independent production storage/lazy-materialization entrypoint supplied", + "C": "unavailable: no independent production compression-algorithm entrypoint supplied", + } + entrypoint_ok = ( + args.factor_entrypoint is not None and args.factor_entrypoint.is_file() + ) + plan = { + "status": "blocked" if not entrypoint_ok else "requires_factor_contract", + "experiment": "real-2^3-factorial", + "factors": factors, + "combinations": ["".join(bits) for bits in itertools.product("0A", repeat=0)], + "baseline": git_identity(args.baseline_repo), + "candidate": git_identity(args.candidate_repo), + "warmup_scope": "sdk_send_only; not warm history", + "reason": "Do not execute or label combinations until A/B/C independently select real production paths.", + } + plan["combinations"] = [ + "".join(bits) for bits in itertools.product("0A", "0B", "0C") + ] + args.output.write_text(json.dumps(plan, indent=2) + "\n", encoding="utf-8") + print(json.dumps({"status": plan["status"], "output": str(args.output)})) + return 0 if plan["status"] == "blocked" else 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/image_memory_bench/baseline_manifest.py b/scripts/image_memory_bench/baseline_manifest.py new file mode 100644 index 0000000000..1710c775cb --- /dev/null +++ b/scripts/image_memory_bench/baseline_manifest.py @@ -0,0 +1,54 @@ +"""Record the exact code and dependency baseline for image experiments.""" + +from __future__ import annotations + +import argparse +import hashlib +import importlib.metadata +import json +import platform +import subprocess +import sys +from pathlib import Path + + +def main() -> int: + """Write a reproducibility manifest without reading application data.""" + parser = argparse.ArgumentParser() + parser.add_argument("output", type=Path) + parser.add_argument("--fixtures", type=Path) + args = parser.parse_args() + commit = subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip() + packages = sorted( + f"{dist.metadata['Name']}=={dist.version}" + for dist in importlib.metadata.distributions() + if dist.metadata.get("Name") + ) + manifest = { + "commit": commit, + "working_tree_diff_sha256": hashlib.sha256( + subprocess.check_output(["git", "diff", "--binary"], cwd=Path.cwd()) + ).hexdigest(), + "python": sys.version, + "platform": platform.platform(), + "packages": packages, + } + if args.fixtures: + files = [] + for path in sorted(args.fixtures.rglob("*")): + if path.is_file() and path.name != "manifest.json": + files.append( + { + "path": str(path.relative_to(args.fixtures)), + "bytes": path.stat().st_size, + "sha256": hashlib.sha256(path.read_bytes()).hexdigest(), + } + ) + manifest["fixtures"] = files + args.output.write_text(json.dumps(manifest, indent=2) + "\n", encoding="utf-8") + print(json.dumps({"commit": commit, "package_count": len(packages)})) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/image_memory_bench/generate_fixtures.py b/scripts/image_memory_bench/generate_fixtures.py new file mode 100644 index 0000000000..74c8cc9a3a --- /dev/null +++ b/scripts/image_memory_bench/generate_fixtures.py @@ -0,0 +1,127 @@ +"""Generate seeded fixtures outside the measured application process.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import random +from pathlib import Path + +from PIL import Image, ImageDraw + + +def generate(output: Path, seed: int) -> list[dict]: + """Create ten inputs for one paired experiment seed. + + Args: + output: New directory for this seed's fixtures. + seed: Stable seed for pixel generation. + + Returns: + On-disk hashes and decoded image metadata. + """ + output.mkdir(parents=True, exist_ok=False) + rng = random.Random(seed) + records = [] + for kind in ("png", "jpeg", "webp", "gif", "screenshot"): + for stress in (False, True): + name = f"{kind}-{'stress' if stress else 'ordinary'}" + options = {} + if kind in ("png", "webp"): + size = (1280, 1280) if stress else (320, 240) + image = ( + Image.frombytes("RGBA", size, rng.randbytes(size[0] * size[1] * 4)) + if stress + else Image.new("RGBA", size, (40, 120, 200, 127)) + ) + image_format = kind.upper() + options = {"lossless": stress} if kind == "webp" else {} + elif kind == "jpeg": + size = (4000, 3000) if stress else (640, 480) + image = Image.frombytes( + "RGB", size, rng.randbytes(size[0] * size[1] * 3) + ) + image_format = "JPEG" + options = {"quality": 95 if stress else 75} + if stress: + exif = image.getexif() + exif[274] = 6 + options["exif"] = exif + elif kind == "gif": + image = Image.new("RGB", (320, 240), (seed % 255, 0, 255)) + image_format = "GIF" + frames = [] + for frame in range(24 if stress else 3): + extra = Image.new("RGB", image.size, (frame * 9, 80, 20)) + ImageDraw.Draw(extra).rectangle( + (frame * 5, 10, frame * 5 + 30, 50), fill="white" + ) + frames.append(extra) + options = { + "save_all": True, + "append_images": frames, + "duration": 80, + "loop": 0, + } + else: + image = Image.new( + "RGB", (3840, 2160) if stress else (1920, 1080), "#202124" + ) + image_format = "PNG" + draw = ImageDraw.Draw(image) + for row in range(image.height // 20): + draw.text( + (100, row * 20), + f"coordinate (100, {row * 20}) seed={seed} command --verbose " + * 4, + fill="#e8eaed", + ) + draw.rectangle((20, 20, 70, 70), outline="red", width=3) + path = output / f"{name}.{image_format.lower()}" + try: + image.save(path, image_format, **options) + finally: + image.close() + for frame_image in options.get("append_images", []): + frame_image.close() + # Reopen the saved bytes: in-memory Image.format and frame counts + # do not describe what the encoder actually persisted. + with Image.open(path) as saved: + record = { + "path": path.name, + "kind": kind, + "stress": stress, + "seed": seed, + "format": saved.format, + "width": saved.width, + "height": saved.height, + "mode": saved.mode, + "frames": getattr(saved, "n_frames", 1), + "orientation": saved.getexif().get(274, 1), + "has_alpha": saved.mode in ("RGBA", "LA") + or "transparency" in saved.info, + "bytes": path.stat().st_size, + "sha256": hashlib.sha256(path.read_bytes()).hexdigest(), + } + records.append(record) + (output / "manifest.json").write_text( + json.dumps(records, indent=2) + "\n", encoding="utf-8" + ) + return records + + +def main() -> int: + """Generate three paired seeds unless explicitly overridden.""" + parser = argparse.ArgumentParser() + parser.add_argument("output", type=Path) + parser.add_argument("--seeds", nargs="+", type=int, default=[7, 19, 43]) + args = parser.parse_args() + for seed in args.seeds: + records = generate(args.output / str(seed), seed) + print(json.dumps({"seed": seed, "count": len(records)})) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/image_memory_bench/history_benchmark.py b/scripts/image_memory_bench/history_benchmark.py new file mode 100644 index 0000000000..0e9563bf6b --- /dev/null +++ b/scripts/image_memory_bench/history_benchmark.py @@ -0,0 +1,461 @@ +"""Measure SQLite history loading and real SDK request bodies for B-only.""" + +# The compact benchmark runner keeps subprocess orchestration visibly linear. +# ruff: noqa: E701, E702 + +from __future__ import annotations + +import argparse +import asyncio +import base64 +import csv +import hashlib +import json +import mimetypes +import os +import subprocess +import sys +import tempfile +import threading +import time +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path + +import psutil + + +async def prepare_database( + db_path: Path, + media_root: Path, + fixture: Path, + images: int, + distinct: bool, + mode: str, +) -> None: + """Create one SQLite input outside the measured child.""" + sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + from astrbot.core.db.sqlite import SQLiteDatabase + from astrbot.core.utils.image_media_store import ( + ImageMediaStore, + persist_inline_image_refs, + ) + + rows = json.loads(fixture.read_text()) + paths = [fixture.parent / row["path"] for row in rows] + history = [] + raw = [] + for index in range(images): + path = paths[index % len(paths)] if distinct else paths[0] + # Distinct workloads use separately generated, valid image files listed + # by the fixture manifest. Never mutate encoded bytes into invalid images. + data = path.read_bytes() + mime = mimetypes.guess_type(path.name)[0] or "image/png" + raw.append( + { + "type": "image_url", + "image_url": { + "url": f"data:{mime};base64,{base64.b64encode(data).decode()}", + "detail": "high", + }, + } + ) + for index, part in enumerate(raw): + history.extend( + [ + {"role": "user", "content": [part]}, + {"role": "assistant", "content": f"ack {index}"}, + ] + ) + history.extend( + [ + {"role": "user", "content": "final text"}, + {"role": "assistant", "content": "final response"}, + ] + ) + if mode == "reference": + history = persist_inline_image_refs(history, ImageMediaStore(media_root)) + db = SQLiteDatabase(str(db_path)) + await db.initialize() + record = await db.create_conversation( + user_id="bench", platform_id="bench", content=history, title="benchmark" + ) + (db_path.parent / "conversation-id").write_text(record.conversation_id) + await db.engine.dispose() + + +def count_images(value: object) -> int: + """Count data image URLs in a decoded SDK body.""" + if isinstance(value, dict): + return sum(count_images(item) for item in value.values()) + if isinstance(value, list): + return sum(count_images(item) for item in value) + return int(isinstance(value, str) and value.startswith("data:image/")) + + +async def run_child(args: argparse.Namespace) -> dict: + """Read through SQLite and ConversationManager before selecting context.""" + sys.path.insert(0, str(args.repo.resolve())) + from openai import AsyncOpenAI + + from astrbot.core.agent.context.config import ContextConfig + from astrbot.core.agent.context.manager import ContextManager + from astrbot.core.agent.message import Message + from astrbot.core.conversation_mgr import ConversationManager + from astrbot.core.db.sqlite import SQLiteDatabase + from astrbot.core.utils.image_media_store import ( + ImageMediaStore, + materialize_image_media_refs, + ) + + phases = {} + + def mark(phase: str) -> None: + with args.marker.open("a", encoding="utf-8") as marker: + marker.write( + json.dumps({"phase": phase, "monotonic": time.monotonic()}) + "\n" + ) + + mark("db_and_convert") + started = time.perf_counter() + db = SQLiteDatabase(str(args.database)) + await db.initialize() + conversation = await ConversationManager(db).get_conversation( + "bench", args.conversation_id + ) + phases["db_and_convert_ms"] = (time.perf_counter() - started) * 1000 + mark("json_load_bind") + started = time.perf_counter() + history = json.loads(conversation.history) + messages = [Message.model_validate(item) for item in history] + phases["json_load_bind_ms"] = (time.perf_counter() - started) * 1000 + mark("selection") + started = time.perf_counter() + selected = await ContextManager( + ContextConfig(enforce_max_turns=args.keep_turns) + ).process(messages) + phases["selection_ms"] = (time.perf_counter() - started) * 1000 + mark("materialize") + selected = await materialize_image_media_refs( + selected, ImageMediaStore(args.media_root) + ) + payload = [message.model_dump() for message in selected] + mark("sdk") + async with AsyncOpenAI( + api_key="local", base_url=args.endpoint, max_retries=0 + ) as client: + for _ in range(args.warmup): + await client.chat.completions.create(model="local", messages=payload) + await client.chat.completions.create(model="local", messages=payload) + await db.engine.dispose() + return { + "phases_ms": phases, + "selected_messages": len(payload), + "selected_images": count_images(payload), + } + + +def main() -> int: + """Prepare paired databases and monitor fresh children.""" + parser = argparse.ArgumentParser() + parser.add_argument("fixture", type=Path) + parser.add_argument("output", type=Path) + parser.add_argument( + "--repo", type=Path, default=Path(__file__).resolve().parents[2] + ) + parser.add_argument("--images", type=int, default=50) + parser.add_argument("--keep-turns", type=int, default=25) + parser.add_argument("--repeats", type=int, default=10) + parser.add_argument("--cold-repeats", type=int) + parser.add_argument("--warm-repeats", type=int) + parser.add_argument("--warmup", type=int, default=3) + parser.add_argument("--distinct", action="store_true") + parser.add_argument("--csv", type=Path) + parser.add_argument("--rss-jsonl", type=Path) + parser.add_argument("--child", action="store_true") + parser.add_argument("--database", type=Path) + parser.add_argument("--media-root", type=Path) + parser.add_argument("--history-mode", choices=("inline", "reference")) + parser.add_argument("--conversation-id", default="") + parser.add_argument("--endpoint") + parser.add_argument("--body-file", type=Path) + parser.add_argument("--marker", type=Path) + args = parser.parse_args() + cold_repeats = args.cold_repeats if args.cold_repeats is not None else args.repeats + warm_repeats = args.warm_repeats if args.warm_repeats is not None else args.repeats + if min(args.repeats, cold_repeats, warm_repeats, args.images, args.keep_turns) < 1: + parser.error("repetition, image, and context-window values must be positive") + fixture_rows = json.loads(args.fixture.read_text()) + max_fixture_bytes = max( + ( + row.get("bytes") or row.get("prepared_bytes") or row.get("source_bytes", 0) + for row in fixture_rows + ), + default=0, + ) + budget_class = ( + "over-4MiB-source-stress" + if 4 * ((max_fixture_bytes + 2) // 3) > 4 * 1024 * 1024 + else "representative-within-4MiB-source" + ) + if args.child: + if ( + args.database is None + or args.media_root is None + or args.endpoint is None + or args.body_file is None + or args.marker is None + or args.history_mode is None + ): + parser.error( + "child requires database, media root, endpoint, body file, marker, and history mode" + ) + result = asyncio.run(run_child(args)) + if args.body_file.exists(): + metadata = json.loads(args.body_file.read_text()) + result.update( + { + "http_json_bytes": metadata["bytes"], + "http_body_sha256": metadata["sha256"], + "http_images": metadata["images"], + } + ) + else: + result.update( + { + "http_json_bytes": None, + "http_body_sha256": None, + "http_images": 0, + } + ) + if not result["selected_images"] or not result["http_images"]: + raise RuntimeError("actual SDK body contains no images") + args.output.write_text(json.dumps(result) + "\n") + return 0 + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text("", encoding="utf-8") + csv_file = args.csv.open("w", newline="", encoding="utf-8") if args.csv else None + writer = None + received_body = {"path": None} + rss_path = args.output.with_name(args.output.stem + ".rss.jsonl") + if args.rss_jsonl: + rss_path = args.rss_jsonl + rss_path.parent.mkdir(parents=True, exist_ok=True) + rss_file = rss_path.open("w", encoding="utf-8") + + class Handler(BaseHTTPRequestHandler): + def do_POST(self): + data = self.rfile.read(int(self.headers.get("Content-Length", "0"))) + body = json.loads(data) + Path(received_body["path"]).write_text( + json.dumps( + { + "bytes": len(data), + "sha256": hashlib.sha256(data).hexdigest(), + "images": count_images(body), + } + ) + ) + response = b'{"choices":[{"message":{"content":"ok"}}]}' + self.send_response(200) + self.send_header("Content-Length", str(len(response))) + self.end_headers() + self.wfile.write(response) + + def log_message(self, *_args): + pass + + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + threading.Thread(target=server.serve_forever, daemon=True).start() + endpoint = f"http://127.0.0.1:{server.server_port}/v1" + try: + for mode in ("inline", "reference"): + for repeat in range(cold_repeats + warm_repeats): + with tempfile.TemporaryDirectory( + prefix="astrbot-history-input-" + ) as temp: + root = Path(temp) + db_path = root / f"{mode}.db" + media = root / mode + asyncio.run( + prepare_database( + db_path, + media, + args.fixture, + args.images, + args.distinct, + mode, + ) + ) + conversation_id = (db_path.parent / "conversation-id").read_text() + child_output = root / "child.jsonl" + body_file = root / "http-metadata.json" + received_body["path"] = body_file + warmup = 0 if repeat < cold_repeats else args.warmup + command = [ + sys.executable, + str(Path(__file__).resolve()), + str(args.fixture), + str(child_output), + "--child", + "--repo", + str(args.repo), + "--database", + str(db_path), + "--media-root", + str(media), + "--conversation-id", + conversation_id, + "--keep-turns", + str(args.keep_turns), + "--history-mode", + mode, + "--warmup", + str(warmup), + "--endpoint", + endpoint, + "--body-file", + str(body_file), + "--marker", + str(root / "marker.jsonl"), + ] + proc = subprocess.Popen( + command, + cwd=args.repo, + env={**os.environ, "ASTRBOT_ROOT": str(root)}, + ) + monitor = psutil.Process(proc.pid) + samples = [] + started = time.monotonic() + limit = min(2 * 1024**3, psutil.virtual_memory().available // 4) + stop_reason = None + while proc.poll() is None: + try: + now = time.monotonic() + rss = monitor.memory_info().rss + private_bytes = None + try: + private_bytes = monitor.memory_full_info().uss + except (AttributeError, psutil.AccessDenied): + pass + samples.append( + { + "seconds": now - started, + "monotonic": now, + "rss": rss, + "private_bytes": private_bytes, + } + ) + if now - started > 120: + stop_reason = "timeout" + elif rss > limit: + stop_reason = "memory_limit" + if stop_reason: + proc.kill() + break + except psutil.NoSuchProcess: + break + time.sleep(0.01) + proc.wait() + if child_output.exists(): + result = json.loads(child_output.read_text()) + else: + result = { + "status": "error", + "error": stop_reason or "child_failed", + } + marker_file = root / "marker.jsonl" + markers = ( + [ + json.loads(line) + for line in marker_file.read_text().splitlines() + ] + if marker_file.exists() + else [] + ) + for sample in samples: + eligible = [ + marker + for marker in markers + if marker["monotonic"] <= sample["monotonic"] + ] + sample["phase"] = ( + max(eligible, key=lambda marker: marker["monotonic"])[ + "phase" + ] + if eligible + else "pre-start" + ) + result.update( + { + "mode": mode, + "repeat": repeat, + "cold": repeat < cold_repeats, + "warmup_mode": "sdk_warmup" if warmup else "cold", + "warmup_scope": "sdk_send_only", + "workload": f"sqlite-history-{args.images}-{'distinct' if args.distinct else 'same'}", + "budget_class": budget_class, + "rss_high_water_bytes": max( + (sample["rss"] for sample in samples), default=None + ), + "private_high_water_bytes": max( + ( + sample["private_bytes"] + for sample in samples + if sample["private_bytes"] is not None + ), + default=None, + ), + "stop_reason": stop_reason, + "rss_sample_count": len(samples), + "rss_jsonl": str(rss_path), + "returncode": proc.returncode, + "status": ( + "ok" + if proc.returncode == 0 and stop_reason is None + else "error" + ), + } + ) + for sample in samples: + rss_file.write( + json.dumps( + { + "mode": mode, + "repeat": repeat, + **sample, + } + ) + + "\n" + ) + rss_file.flush() + with args.output.open("a", encoding="utf-8") as stream: + stream.write(json.dumps(result) + "\n") + if csv_file: + fields = [ + "mode", + "repeat", + "cold", + "workload", + "rss_high_water_bytes", + "private_high_water_bytes", + "rss_sample_count", + "status", + "returncode", + "http_json_bytes", + "http_body_sha256", + "http_images", + ] + if writer is None: + writer = csv.DictWriter(csv_file, fieldnames=fields) + writer.writeheader() + writer.writerow({field: result.get(field) for field in fields}) + finally: + server.shutdown() + rss_file.close() + if csv_file: + csv_file.close() + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/image_memory_bench/lifecycle_workloads.py b/scripts/image_memory_bench/lifecycle_workloads.py new file mode 100644 index 0000000000..75c80846e0 --- /dev/null +++ b/scripts/image_memory_bench/lifecycle_workloads.py @@ -0,0 +1,371 @@ +"""Run real history lifecycle workloads in an isolated, measured child process.""" + +from __future__ import annotations + +import argparse +import asyncio +import base64 +import hashlib +import json +import subprocess +import sys +import tempfile +import threading +import time +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from typing import Any + +import psutil + +REQUESTS = 200 +SESSIONS = 4 +WINDOW_TURNS = 25 + + +def _count_images(value: Any) -> int: + if isinstance(value, dict): + return sum(_count_images(item) for item in value.values()) + if isinstance(value, list): + return sum(_count_images(item) for item in value) + return int(isinstance(value, str) and value.startswith("data:image/")) + + +def _image_hashes(value: Any) -> list[str]: + if isinstance(value, dict): + if value.get("type") == "image_url" and isinstance( + value.get("image_url"), dict + ): + encoded = value["image_url"].get("url", "").split(",", 1)[-1] + try: + return [hashlib.sha256(base64.b64decode(encoded)).hexdigest()] + except Exception: + return [] + return [item for child in value.values() for item in _image_hashes(child)] + if isinstance(value, list): + return [item for child in value for item in _image_hashes(child)] + return [] + + +async def _child(args: argparse.Namespace) -> dict[str, Any]: + """Execute the production database, context, materialization, and SDK path.""" + sys.path.insert(0, str(args.repo.resolve())) + from openai import AsyncOpenAI + + from astrbot.core.agent.context.config import ContextConfig + from astrbot.core.agent.context.manager import ContextManager + from astrbot.core.agent.message import Message + from astrbot.core.conversation_mgr import ConversationManager + from astrbot.core.db.sqlite import SQLiteDatabase + from astrbot.core.utils.image_media_store import ( + ImageMediaStore, + materialize_image_media_refs, + ) + + async def load_payload( + database: SQLiteDatabase, conversation_id: str + ) -> tuple[list, str, list]: + conversation = await ConversationManager(database).get_conversation( + "bench", conversation_id + ) + messages = [ + Message.model_validate(item) for item in json.loads(conversation.history) + ] + selected = await ContextManager( + ContextConfig(enforce_max_turns=args.window_turns) + ).process(messages) + selected = await materialize_image_media_refs( + selected, ImageMediaStore(args.media_root) + ) + payload = [message.model_dump() for message in selected] + semantic_hash = hashlib.sha256( + json.dumps(payload, sort_keys=True, ensure_ascii=False).encode() + ).hexdigest() + return payload, semantic_hash, messages + + async def session( + client: AsyncOpenAI, session_id: str, request_count: int + ) -> dict[str, Any]: + database = SQLiteDatabase(str(args.database)) + await database.initialize() + semantic_hashes: list[str] = [] + image_counts: list[int] = [] + image_hash_orders: list[list[str]] = [] + for turn in range(request_count): + _, _, messages = await load_payload(database, session_id) + image_refs = [ + part.model_dump() if hasattr(part, "model_dump") else part + for message in messages + for part in ( + message.content if isinstance(message.content, list) else [] + ) + if (part.model_dump() if hasattr(part, "model_dump") else part).get( + "type" + ) + == "image_media_ref" + ] + if not image_refs: + raise RuntimeError( + "lifecycle fixture has no persisted image references" + ) + messages.extend( + [ + Message(role="user", content=[image_refs[turn % len(image_refs)]]), + Message(role="assistant", content=f"ack {turn}"), + ] + ) + await ConversationManager(database).update_conversation( + "bench", + session_id, + history=[message.model_dump() for message in messages], + ) + payload, semantic_hash, _ = await load_payload(database, session_id) + semantic_hashes.append(semantic_hash) + image_counts.append(_count_images(payload)) + image_hash_orders.append(_image_hashes(payload)) + if not image_hash_orders[-1] or image_counts[-1] > args.window_turns: + raise RuntimeError("active image window is empty or unbounded") + await client.chat.completions.create(model="local", messages=payload) + del payload + await database.engine.dispose() + return { + "session": session_id, + "requests": request_count, + "image_counts": image_counts, + "image_hash_orders": image_hash_orders, + "semantic_hashes": semantic_hashes, + } + + if args.restart_check: + checks = [] + database = SQLiteDatabase(str(args.database)) + await database.initialize() + for session_id in args.conversation_ids: + payload, semantic_hash, _ = await load_payload(database, session_id) + checks.append( + { + "session": session_id, + "images": _count_images(payload), + "semantic_hash": semantic_hash, + "image_hash_order": _image_hashes(payload), + } + ) + await database.engine.dispose() + return {"restart_checks": checks} + + async with AsyncOpenAI( + api_key="local", base_url=args.endpoint, max_retries=0 + ) as client: + session_count = len(args.conversation_ids) + base_requests, remainder = divmod(args.request_count, session_count) + request_counts = [ + base_requests + (index < remainder) for index in range(session_count) + ] + results = await asyncio.gather( + *( + session(client, cid, request_count) + for cid, request_count in zip(args.conversation_ids, request_counts) + ) + ) + return { + "sessions": results, + "total_requests": sum(request_counts), + "request_count_per_session": request_counts, + "window_turns": args.window_turns, + } + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("output", type=Path) + parser.add_argument( + "--repo", type=Path, default=Path(__file__).resolve().parents[2] + ) + parser.add_argument("--child", action="store_true") + parser.add_argument("--database", type=Path) + parser.add_argument("--media-root", type=Path) + parser.add_argument( + "--conversation-id", dest="conversation_ids", action="append", default=[] + ) + parser.add_argument("--endpoint") + parser.add_argument("--rss-jsonl", type=Path) + parser.add_argument("--restart-check", action="store_true") + parser.add_argument("--session-count", type=int, choices=(1, 4), default=4) + parser.add_argument("--request-count", type=int, default=REQUESTS) + parser.add_argument("--window-turns", type=int, default=WINDOW_TURNS) + args = parser.parse_args() + if args.request_count < 1 or args.window_turns < 1: + parser.error("request count and window turns must be positive") + if args.child: + if ( + args.database is None + or args.media_root is None + or args.endpoint is None + or len(args.conversation_ids) != args.session_count + ): + parser.error( + "child requires database, media root, endpoint, and the requested conversation IDs" + ) + result = asyncio.run(_child(args)) + args.output.write_text(json.dumps(result) + "\n", encoding="utf-8") + return 0 + if len(args.conversation_ids) != args.session_count: + parser.error( + f"exactly {args.session_count} --conversation-id values are required" + ) + if args.database is None or args.media_root is None: + parser.error("--database and --media-root are required") + + requests: list[dict[str, Any]] = [] + + class Sink(BaseHTTPRequestHandler): + def do_POST(self) -> None: + size = int(self.headers.get("Content-Length", "0")) + body = self.rfile.read(size) + requests.append( + { + "bytes": len(body), + "sha256": hashlib.sha256(body).hexdigest(), + "images": _count_images(json.loads(body)), + } + ) + response = b'{"choices":[{"message":{"content":"ok"}}]}' + self.send_response(200) + self.send_header("Content-Length", str(len(response))) + self.end_headers() + self.wfile.write(response) + + def log_message(self, *_args: Any) -> None: + pass + + server = ThreadingHTTPServer(("127.0.0.1", 0), Sink) + threading.Thread(target=server.serve_forever, daemon=True).start() + with tempfile.TemporaryDirectory(prefix="astrbot-lifecycle-") as run_dir: + child_output = Path(run_dir) / "child.json" + rss_path = args.rss_jsonl or args.output.with_name( + args.output.stem + ".rss.jsonl" + ) + command = [ + sys.executable, + str(Path(__file__).resolve()), + str(child_output), + "--child", + "--repo", + str(args.repo.resolve()), + "--database", + str(args.database.resolve()), + "--media-root", + str(args.media_root.resolve()), + "--endpoint", + f"http://127.0.0.1:{server.server_port}/v1", + "--session-count", + str(args.session_count), + "--request-count", + str(args.request_count), + "--window-turns", + str(args.window_turns), + ] + for conversation_id in args.conversation_ids: + command.extend(["--conversation-id", conversation_id]) + started = time.monotonic() + process = subprocess.Popen(command, cwd=args.repo) + child = psutil.Process(process.pid) + stop_reason = None + with rss_path.open("w", encoding="utf-8") as rss: + while process.poll() is None: + current = child.memory_info().rss + private_bytes = None + try: + private_bytes = child.memory_full_info().uss + except (AttributeError, psutil.AccessDenied): + pass + rss.write( + json.dumps( + { + "seconds": time.monotonic() - started, + "rss": current, + "private_bytes": private_bytes, + } + ) + + "\n" + ) + rss.flush() + if ( + current > min(2 * 1024**3, psutil.virtual_memory().available // 4) + or time.monotonic() - started > 120 + ): + stop_reason = ( + "memory_limit" + if current + > min(2 * 1024**3, psutil.virtual_memory().available // 4) + else "timeout" + ) + process.kill() + break + time.sleep(0.01) + return_code = process.wait() + server.shutdown() + result = ( + json.loads(child_output.read_text()) + if child_output.exists() + else {"error": "child failed"} + ) + validation_errors = [] + restart_result = {"skipped": stop_reason or "child_failed"} + if stop_reason is None and return_code == 0: + restart_output = Path(run_dir) / "restart.json" + restart_command = command.copy() + restart_command[2] = str(restart_output) + restart_command.append("--restart-check") + restart_process = subprocess.run( + restart_command, cwd=args.repo, check=False + ) + restart_result = ( + json.loads(restart_output.read_text()) + if restart_output.exists() + else {"error": "restart child failed"} + ) + if restart_result.get("restart_checks"): + expected = { + item["session"]: item for item in result.get("sessions", []) + } + for check in restart_result["restart_checks"]: + prior = expected.get(check["session"]) + if ( + prior is None + or check["images"] != prior["image_counts"][-1] + or check["image_hash_order"] != prior["image_hash_orders"][-1] + ): + validation_errors.append( + "fresh-process restart changed the active image window" + ) + if restart_process.returncode != 0: + return_code = restart_process.returncode + result.update( + { + "returncode": return_code, + "stop_reason": stop_reason, + "request_count": len(requests), + "expected_request_count": args.request_count, + "requests": requests, + "rss_jsonl": str(rss_path), + "restart": restart_result, + } + ) + if any( + request["images"] <= 0 or request["images"] > args.window_turns + for request in requests + ): + validation_errors.append("sink observed an invalid active image window") + if len(requests) != args.request_count: + validation_errors.append("sink did not observe the expected request count") + if validation_errors: + result["validation_errors"] = validation_errors + return_code = return_code or 1 + result["returncode"] = return_code + args.output.write_text(json.dumps(result) + "\n", encoding="utf-8") + return return_code + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/image_memory_bench/measure.py b/scripts/image_memory_bench/measure.py new file mode 100644 index 0000000000..74c4bdb26e --- /dev/null +++ b/scripts/image_memory_bench/measure.py @@ -0,0 +1,399 @@ +"""Measure real image preparation and SDK requests in isolated processes.""" + +from __future__ import annotations + +import argparse +import asyncio +import base64 +import csv +import hashlib +import json +import os +import subprocess +import sys +import tempfile +import threading +import time +import tracemalloc +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path + +import psutil + + +def _base64_payload_metrics(encoded: str) -> tuple[int, str]: + """Hash an encoded payload in chunks without retaining decoded image bytes. + + Args: + encoded: Base64 payload without the data-URI prefix. + + Returns: + The decoded byte count and SHA-256 digest. + """ + digest = hashlib.sha256() + decoded_bytes = 0 + remainder = "" + chunk_size = 4 * 1024 * 1024 + for start in range(0, len(encoded), chunk_size): + chunk = remainder + encoded[start : start + chunk_size] + usable = len(chunk) - len(chunk) % 4 + if usable: + decoded = base64.b64decode(chunk[:usable]) + digest.update(decoded) + decoded_bytes += len(decoded) + remainder = chunk[usable:] + if remainder: + decoded = base64.b64decode(remainder) + digest.update(decoded) + decoded_bytes += len(decoded) + return decoded_bytes, digest.hexdigest() + + +async def run_child(args: argparse.Namespace) -> dict: + """Execute production preparation and request assembly in the chosen checkout. + + Args: + args: Explicit source, checkout, local server and diagnostic options. + + Returns: + Sizes, phase durations and resource counters without image content. + """ + sys.path.insert(0, str(args.repo.resolve())) + from openai import AsyncOpenAI + + from astrbot.core.provider.entities import ProviderRequest + + if args.trace_python: + tracemalloc.start() + process = psutil.Process() + baseline_rss = process.memory_info().rss + phases: dict[str, float] = {} + stage = "startup" + phase_marks = [(time.monotonic(), stage)] + failure: dict[str, str] | None = None + source_bytes = args.image.stat().st_size + source_digest = hashlib.sha256() + with args.image.open("rb") as source_stream: + for chunk in iter(lambda: source_stream.read(1024 * 1024), b""): + source_digest.update(chunk) + source_sha256 = source_digest.hexdigest() + prepared_bytes: int | None = None + prepared_sha256: str | None = None + started = time.perf_counter() + with tempfile.TemporaryDirectory(prefix="astrbot-image-runtime-") as runtime: + os.environ["ASTRBOT_ROOT"] = runtime + before_files = set(Path(runtime).rglob("*")) + try: + stage = "provider_prepare_and_assemble" + phase_marks.append((time.monotonic(), stage)) + tick = time.perf_counter() + request = ProviderRequest( + prompt="Describe this image.", + image_urls=[str(args.image.resolve())], + ) + context = await request.assemble_context() + phases["provider_prepare_and_assemble_ms"] = ( + time.perf_counter() - tick + ) * 1000 + image_part = next( + part + for part in context.get("content", []) + if part.get("type") == "image_url" + ) + data_url = image_part["image_url"]["url"] + stage = "benchmark_payload_verification" + phase_marks.append((time.monotonic(), stage)) + prepared_bytes, prepared_sha256 = _base64_payload_metrics( + data_url.split(",", 1)[1] + ) + stage = "sdk_send" + phase_marks.append((time.monotonic(), stage)) + async with AsyncOpenAI( + api_key="local-experiment-only", + base_url=args.endpoint, + max_retries=0, + ) as client: + for _ in range(args.warmup): + await client.chat.completions.create( + model="local", messages=[context] + ) + tick = time.perf_counter() + for _ in range(args.repeat): + await client.chat.completions.create( + model="local", messages=[context] + ) + phases["sdk_send_ms"] = (time.perf_counter() - tick) * 1000 + del context, request + except Exception as exc: + failure = {"stage": stage, "error_type": type(exc).__name__} + finally: + residual = [ + path + for path in Path(runtime).rglob("*") + if path.is_file() and path not in before_files + ] + result = { + "source_bytes": source_bytes, + "prepared_bytes": prepared_bytes, + "source_sha256": source_sha256, + "prepared_sha256": prepared_sha256, + "preserved_source_bytes": ( + prepared_sha256 == source_sha256 + if prepared_sha256 is not None + else None + ), + "base64_bytes": ( + 4 * ((prepared_bytes + 2) // 3) + if prepared_bytes is not None + else None + ), + "phase_ms": phases, + "phase_marks": phase_marks, + "last_stage": stage, + "failure": failure, + "baseline_rss": baseline_rss, + "post_request_rss": process.memory_info().rss, + "temporary_files_remaining": len(residual), + "temporary_bytes_remaining": sum( + path.stat().st_size for path in residual + ), + "elapsed_ms": (time.perf_counter() - started) * 1000, + } + if args.trace_python: + result["python_current_bytes"], result["python_peak_bytes"] = ( + tracemalloc.get_traced_memory() + ) + tracemalloc.stop() + if sys.platform != "win32": + import resource + + peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + result["os_peak_rss_bytes"] = ( + peak if sys.platform == "darwin" else peak * 1024 + ) + return result + + +def main() -> int: + """Run an isolated request and record measurements or a bounded failure.""" + parser = argparse.ArgumentParser() + parser.add_argument("image", type=Path) + parser.add_argument("output", type=Path) + parser.add_argument( + "--repo", type=Path, default=Path(__file__).resolve().parents[2] + ) + parser.add_argument("--trace-python", action="store_true") + parser.add_argument("--child", action="store_true") + parser.add_argument("--endpoint") + parser.add_argument("--repeat", type=int, default=1) + parser.add_argument("--warmup", type=int, default=0) + parser.add_argument("--csv") + parser.add_argument("--workload", default="single-image") + args = parser.parse_args() + if args.child: + try: + result = asyncio.run(run_child(args)) + result["status"] = "error" if result.get("failure") else "ok" + except Exception as error: + result = {"status": "error", "error_type": type(error).__name__} + args.output.write_text(json.dumps(result), encoding="utf-8") + return 0 if result["status"] == "ok" else 1 + + # The server measures actual SDK bytes without storing the request body. + request_sizes = [] + + class Handler(BaseHTTPRequestHandler): + def do_POST(self): + length = int(self.headers.get("Content-Length", "0")) + remaining = length + while remaining: + chunk = self.rfile.read(min(65536, remaining)) + if not chunk: + break + remaining -= len(chunk) + request_sizes.append(length - remaining) + body = b'{"id":"local","object":"chat.completion","created":0,"model":"local","choices":[{"index":0,"message":{"role":"assistant","content":"fixture response"},"finish_reason":"stop"}],"usage":{"prompt_tokens":100,"completion_tokens":2,"total_tokens":102}}' + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, *args): + pass + + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + server_thread = threading.Thread(target=server.serve_forever, daemon=True) + server_thread.start() + limit = min(2 * 1024**3, psutil.virtual_memory().available // 4) + started = time.monotonic() + samples = [] + stop_reason = None + try: + with tempfile.TemporaryDirectory(prefix="astrbot-image-measure-") as run_dir: + child_result = Path(run_dir) / "child.json" + command = [ + sys.executable, + str(Path(__file__).resolve()), + str(args.image.resolve()), + str(child_result), + "--child", + "--repo", + str(args.repo.resolve()), + "--endpoint", + f"http://127.0.0.1:{server.server_port}/v1", + ] + if args.repeat != 1: + command.extend(["--repeat", str(args.repeat)]) + if args.warmup: + command.extend(["--warmup", str(args.warmup)]) + command.extend(["--workload", args.workload]) + if args.trace_python: + command.append("--trace-python") + env = os.environ.copy() + env["ASTRBOT_ROOT"] = str(Path(run_dir) / "runtime") + with ( + (Path(run_dir) / "stdout").open("wb") as out, + (Path(run_dir) / "stderr").open("wb") as err, + ): + proc = subprocess.Popen( + command, cwd=args.repo, env=env, stdout=out, stderr=err + ) + process = psutil.Process(proc.pid) + try: + while proc.poll() is None: + try: + children = process.children(recursive=True) + rss = process.memory_info().rss + descendants_rss = sum( + p.memory_info().rss for p in children if p.is_running() + ) + info = process.memory_info() + samples.append( + { + "seconds": time.monotonic() - started, + "rss": rss, + "descendants_rss": descendants_rss, + "private_bytes": getattr(info, "private", None), + } + ) + if rss + descendants_rss > limit: + stop_reason = "memory_limit" + elif time.monotonic() - started > 120: + stop_reason = "timeout" + if stop_reason: + for child in children: + child.kill() + proc.kill() + break + except psutil.NoSuchProcess: + break + time.sleep(0.01) + finally: + if proc.poll() is None: + proc.kill() + proc.wait() + result = ( + json.loads(child_result.read_text()) + if child_result.exists() + else {"status": "error"} + ) + phase_peaks = {} + for sample in samples: + phase = "process_startup" + for timestamp, name in result.get("phase_marks", []): + if timestamp > started + sample["seconds"]: + break + phase = name + sample["phase"] = phase + phase_peaks[phase] = max(phase_peaks.get(phase, 0), sample["rss"]) + result.update( + { + "status": stop_reason or result["status"], + "workload": args.workload, + "phase_rss_high_water_bytes": phase_peaks, + "returncode": proc.returncode, + "rss_high_water_bytes": max( + (s["rss"] for s in samples), default=None + ), + "private_high_water_bytes": max( + ( + s["private_bytes"] + for s in samples + if s["private_bytes"] is not None + ), + default=None, + ), + "descendants_high_water_bytes": max( + (s["descendants_rss"] for s in samples), default=None + ), + "request_json_bytes": request_sizes, + "samples": samples, + "trace_python": args.trace_python, + "code_commit": subprocess.check_output( + ["git", "rev-parse", "HEAD"], cwd=args.repo, text=True + ).strip(), + "working_tree_diff_sha256": hashlib.sha256( + subprocess.check_output( + ["git", "diff", "--binary"], cwd=args.repo + ) + ).hexdigest(), + "benchmark_source_sha256": hashlib.sha256( + Path(__file__).read_bytes() + ).hexdigest(), + "monitor_sample_count": len(samples), + } + ) + with args.output.open("a", encoding="utf-8") as output: + output.write(json.dumps(result) + "\n") + if args.csv: + fields = [ + "status", + "workload", + "repeat", + "warmup", + "rss_high_water_bytes", + "private_high_water_bytes", + "post_request_rss", + "temporary_files_remaining", + "temporary_bytes_remaining", + "request_json_bytes", + ] + csv_path = Path(args.csv) + write_header = not csv_path.exists() + with csv_path.open("a", newline="", encoding="utf-8") as csv_file: + writer = csv.DictWriter(csv_file, fieldnames=fields) + if write_header: + writer.writeheader() + writer.writerow( + { + "status": result.get("status"), + "workload": args.workload, + "repeat": args.repeat, + "warmup": args.warmup, + "rss_high_water_bytes": result.get("rss_high_water_bytes"), + "private_high_water_bytes": result.get( + "private_high_water_bytes" + ), + "post_request_rss": result.get("post_request_rss"), + "temporary_files_remaining": result.get( + "temporary_files_remaining" + ), + "temporary_bytes_remaining": result.get( + "temporary_bytes_remaining" + ), + "request_json_bytes": json.dumps( + result.get("request_json_bytes", []) + ), + } + ) + print(json.dumps({k: v for k, v in result.items() if k != "samples"})) + return 0 if result["status"] == "ok" else 1 + finally: + server.shutdown() + server.server_close() + server_thread.join() + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/image_memory_bench/migrate_history.py b/scripts/image_memory_bench/migrate_history.py new file mode 100644 index 0000000000..d2189d7c95 --- /dev/null +++ b/scripts/image_memory_bench/migrate_history.py @@ -0,0 +1,99 @@ +"""Explicitly externalize or re-inline exported JSONL conversation histories. + +The input is never modified. Inspect the dry run before using --apply. Keep the +input alongside the media directory as the rollback backup. +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + +from astrbot.core.utils.image_media_store import ( + ImageMediaStore, + materialize_image_media_refs, + persist_inline_image_refs, +) + + +async def main() -> int: + """Transform an export into a new file without modifying live histories.""" + parser = argparse.ArgumentParser() + parser.add_argument("input", type=Path) + parser.add_argument("output", type=Path) + parser.add_argument("--media-dir", type=Path, required=True) + parser.add_argument( + "--mode", choices=["externalize", "inline"], default="externalize" + ) + parser.add_argument("--apply", action="store_true") + parser.add_argument("--rollback", action="store_true") + args = parser.parse_args() + if args.input.resolve() == args.output.resolve() or args.output.exists(): + parser.error("Output must be a new file distinct from the rollback input") + if args.rollback: + if not args.apply: + parser.error("--rollback requires --apply") + # The input is the verified original export; rollback writes that exact + # export to a new path, so the original remains the recovery anchor. + args.mode = "inline" + store = ImageMediaStore(args.media_dir) + record_count = 0 + image_count = 0 + output = args.output.open("x", encoding="utf-8") if args.apply else None + try: + with args.input.open(encoding="utf-8") as source: + for line in source: + record = json.loads(line) + raw_history = record.get("history", []) + history = ( + json.loads(raw_history) + if isinstance(raw_history, str) + else raw_history + ) + if not isinstance(history, list): + raise ValueError("History must be a message list") + for message in history: + parts = message.get("content") + if isinstance(parts, list): + image_count += sum( + isinstance(part, dict) + and part.get("type") in {"image_url", "image_media_ref"} + for part in parts + ) + if args.apply: + if args.mode == "externalize" and not args.rollback: + history = persist_inline_image_refs(history, store) + else: + history = await materialize_image_media_refs( + history, store, strict=True + ) + record["history"] = ( + json.dumps(history, ensure_ascii=False) + if isinstance(raw_history, str) + else history + ) + output.write(json.dumps(record, ensure_ascii=False) + "\n") + record_count += 1 + except BaseException: + if output is not None: + output.close() + args.output.unlink(missing_ok=True) + raise + finally: + if output is not None: + output.close() + print( + json.dumps( + {"records": record_count, "images": image_count, "applied": args.apply} + ) + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(asyncio.run(main())) diff --git a/scripts/image_memory_bench/prepare_lifecycle_fixture.py b/scripts/image_memory_bench/prepare_lifecycle_fixture.py new file mode 100644 index 0000000000..b67f9452e5 --- /dev/null +++ b/scripts/image_memory_bench/prepare_lifecycle_fixture.py @@ -0,0 +1,131 @@ +"""Prepare a durable multi-session history fixture before measurement.""" + +from __future__ import annotations + +import argparse +import asyncio +import json +from copy import deepcopy +from pathlib import Path + + +async def prepare_fixture( + manifest_path: Path, + output: Path, + session_count: int, + history_turns: int, + include_stress: bool, +) -> dict: + """Create shared media and persisted conversations outside the measured child. + + Args: + manifest_path: Fixture manifest containing valid prepared image paths. + output: Fresh directory receiving the SQLite database and media objects. + session_count: Number of conversations to create. + history_turns: Number of image-bearing turns per conversation. + include_stress: Whether to include high-entropy stress fixtures. + + Returns: + A JSON-compatible manifest consumed by ``lifecycle_workloads.py``. + + Raises: + FileExistsError: The output already contains a database. + ValueError: The input manifest or requested sizes are invalid. + """ + if session_count < 1 or history_turns < 1: + raise ValueError("session count and history turns must be positive") + rows = json.loads(manifest_path.read_text(encoding="utf-8")) + if not isinstance(rows, list) or not rows: + raise ValueError("fixture manifest must contain at least one image") + if not include_stress: + rows = [row for row in rows if not row.get("stress")] + if not rows: + raise ValueError("fixture manifest must contain at least one image") + output.mkdir(parents=True, exist_ok=True) + database_path = output / "history.db" + if database_path.exists(): + raise FileExistsError(f"refusing to overwrite {database_path}") + + import sys + + sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + from astrbot.core.db.sqlite import SQLiteDatabase + from astrbot.core.utils.image_media_store import ImageMediaStore + + media_root = output / "media" + store = ImageMediaStore(media_root) + refs = [] + for row in rows: + path = manifest_path.parent / row["path"] + refs.append(store.put(path.read_bytes(), detail="high")) + + history = [] + for turn in range(history_turns): + history.extend( + [ + { + "role": "user", + "content": [ + {"type": "text", "text": f"fixture turn {turn}"}, + refs[turn % len(refs)].model_dump(), + ], + }, + {"role": "assistant", "content": f"ack {turn}"}, + ] + ) + + database = SQLiteDatabase(str(database_path)) + await database.initialize() + conversation_ids = [] + for index in range(session_count): + conversation = await database.create_conversation( + user_id="bench", + platform_id="bench", + content=deepcopy(history), + title=f"lifecycle fixture {index}", + ) + conversation_ids.append(conversation.conversation_id) + await database.engine.dispose() + return { + "database": str(database_path.resolve()), + "media_root": str(media_root.resolve()), + "conversation_ids": conversation_ids, + "session_count": session_count, + "history_turns": history_turns, + "media_count": len(refs), + } + + +async def main_async(args: argparse.Namespace) -> int: + """Prepare the fixture and write its machine-readable manifest.""" + result = await prepare_fixture( + args.manifest, + args.output, + args.session_count, + args.history_turns, + args.include_stress, + ) + args.output.mkdir(parents=True, exist_ok=True) + output_manifest = args.output / "lifecycle-manifest.json" + output_manifest.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8") + print(json.dumps(result)) + return 0 + + +def main() -> int: + """Parse fixture preparation arguments.""" + parser = argparse.ArgumentParser() + parser.add_argument("manifest", type=Path) + parser.add_argument("output", type=Path) + parser.add_argument("--session-count", type=int, default=4) + parser.add_argument("--history-turns", type=int, default=50) + parser.add_argument( + "--include-stress", + action="store_true", + help="Include high-entropy stress fixtures in the long-running workload.", + ) + return asyncio.run(main_async(parser.parse_args())) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/manage_image_history.py b/scripts/manage_image_history.py new file mode 100644 index 0000000000..4a3635bba0 --- /dev/null +++ b/scripts/manage_image_history.py @@ -0,0 +1,434 @@ +#!/usr/bin/env python3 +"""Explicit offline migration, repair, rollback, and cleanup for image history.""" + +from __future__ import annotations + +import argparse +import base64 +import hashlib +import json +import shutil +import sqlite3 +import sys +import tempfile +from collections.abc import Iterator +from pathlib import Path +from typing import Any + +from astrbot.core.utils.image_media_store import ImageMediaRef, ImageMediaStore + + +def _digest(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _logical_digest(connection: sqlite3.Connection) -> str: + """Hash every table schema and row without loading the database in memory.""" + digest = hashlib.sha256() + tables = connection.execute( + "SELECT name, sql, type FROM sqlite_master WHERE type IN ('table','index','trigger','view') ORDER BY name" + ) + for name, sql, object_type in tables: + digest.update(json.dumps([name, sql], ensure_ascii=False).encode()) + if object_type != "table": + continue + rows = connection.execute( + f'SELECT * FROM "{name.replace(chr(34), chr(34) * 2)}"' + ) + for row in rows: + digest.update( + json.dumps( + list(row), ensure_ascii=False, default=str, separators=(",", ":") + ).encode() + ) + return digest.hexdigest() + + +def _manifest_media_name(name: str) -> tuple[str, str]: + path = Path(name) + if path.name != name or path.suffix not in {".bin", ".json"}: + raise ValueError("invalid media manifest filename") + media_id = path.stem + if len(media_id) != 64 or any(char not in "0123456789abcdef" for char in media_id): + raise ValueError("invalid media manifest media id") + return media_id, path.suffix + + +def _inline_parts(value: Any) -> Iterator[tuple[dict[str, Any], str, str]]: + if isinstance(value, list): + for item in value: + yield from _inline_parts(item) + elif isinstance(value, dict): + if value.get("type") == "image_url" and isinstance( + value.get("image_url"), dict + ): + url = value["image_url"].get("url") + if isinstance(url, str) and url.startswith("data:image/"): + header, encoded = url.split(",", 1) + yield value, header[5:].split(";", 1)[0], encoded + else: + for child in value.values(): + yield from _inline_parts(child) + + +def _replace_inline(value: Any, store: ImageMediaStore) -> tuple[Any, int, int]: + changed = 0 + bytes_migrated = 0 + + def visit(item: Any) -> Any: + nonlocal changed, bytes_migrated + if isinstance(item, list): + return [visit(child) for child in item] + if not isinstance(item, dict): + return item + if item.get("type") == "image_url" and isinstance(item.get("image_url"), dict): + image = item["image_url"] + url = image.get("url") + if isinstance(url, str) and url.startswith("data:image/"): + header, encoded = url.split(",", 1) + data = base64.b64decode(encoded, validate=True) + ref = store.put( + data, + header[5:].split(";", 1)[0], + image.get("detail"), + image.get("id"), + ) + changed += 1 + bytes_migrated += len(data) + return ref.model_dump() + return {key: visit(child) for key, child in item.items()} + + return visit(value), changed, bytes_migrated + + +def _connect(path: Path) -> sqlite3.Connection: + connection = sqlite3.connect(path) + connection.row_factory = sqlite3.Row + connection.execute("PRAGMA foreign_keys=ON") + return connection + + +def _backup(db: Path, media: Path, destination: Path) -> Path: + """Create a consistent database snapshot and a verified media snapshot.""" + destination.mkdir(parents=True, exist_ok=False) + snapshot = sqlite3.connect(destination / db.name) + readonly = sqlite3.connect(f"file:{db}?mode=ro", uri=True) + try: + readonly.backup(snapshot) + snapshot.commit() + finally: + readonly.close() + snapshot.close() + if media.exists(): + (destination / "media").mkdir() + manifest_media = [] + for path in media.iterdir(): + if path.is_symlink() or not path.is_file(): + raise RuntimeError("media backup refuses symlinks and non-files") + target = destination / "media" / path.name + shutil.copy2(path, target) + manifest_media.append({"name": path.name, "sha256": _digest(target)}) + else: + manifest_media = [] + manifest = { + "database": _digest(destination / db.name), + "media": manifest_media, + } + if (destination / "media").exists(): + manifest["media"] = sorted(manifest_media, key=lambda item: item["name"]) + (destination / "manifest.json").write_text( + json.dumps(manifest, sort_keys=True) + "\n" + ) + return destination + + +def _migrate(args: argparse.Namespace) -> dict[str, int]: + if args.apply and not args.offline: + raise SystemExit("--apply requires --offline: stop all history writes first") + db = Path(args.database).resolve() + media = Path(args.media).resolve() + connection = _connect(db) + lock = args.apply + try: + if lock: + connection.execute("BEGIN IMMEDIATE") + backup = _backup( + db, + media, + Path(args.backup) + if args.backup + else Path( + tempfile.mkdtemp(prefix="astrbot-image-history-backup-parent-") + ) + / "snapshot", + ) + query = ( + "SELECT inner_conversation_id, conversation_id, content FROM conversations" + ) + parameters: tuple[Any, ...] = () + if args.conversation_id: + query += ( + " WHERE conversation_id IN (" + + ",".join("?" for _ in args.conversation_id) + + ")" + ) + parameters = tuple(args.conversation_id) + total = changed = bytes_migrated = 0 + store = ImageMediaStore(media) + for row in connection.execute(query, parameters): + total += 1 + try: + content = ( + json.loads(row["content"]) + if isinstance(row["content"], str) + else row["content"] + ) + converted, count, amount = ( + _replace_inline(content, store) + if args.apply + else (content, sum(1 for _ in _inline_parts(content)), 0) + ) + except Exception as exc: # noqa: BLE001 + raise RuntimeError( + f"conversation {row['conversation_id']} cannot be parsed or verified" + ) from exc + changed += count + bytes_migrated += amount + if args.apply and count: + connection.execute( + "UPDATE conversations SET content=? WHERE inner_conversation_id=?", + ( + json.dumps(converted, ensure_ascii=False), + row["inner_conversation_id"], + ), + ) + if args.apply: + post_migration_digest = _logical_digest(connection) + connection.commit() + manifest_path = backup / "manifest.json" + manifest = json.loads(manifest_path.read_text()) + manifest["post_migration_logical_digest"] = post_migration_digest + manifest_path.write_text(json.dumps(manifest, sort_keys=True) + "\n") + print( + json.dumps( + { + "backup": str(backup), + "records": total, + "images": changed, + "bytes": bytes_migrated, + } + ) + ) + else: + connection.rollback() + print(json.dumps({"dry_run": True, "records": total, "images": changed})) + return {"records": total, "images": changed} + finally: + connection.close() + + +def _cleanup(args: argparse.Namespace) -> None: + if args.apply and not args.offline: + raise SystemExit("cleanup --apply requires --offline") + connection = _connect(Path(args.database).resolve()) + media = Path(args.media).resolve() + try: + connection.execute("BEGIN IMMEDIATE") + referenced: set[str] = set() + for row in connection.execute("SELECT content FROM conversations"): + try: + content = json.loads(row[0]) if isinstance(row[0], str) else row[0] + if not isinstance(content, list): + raise ValueError("history content is not a list") + for item in content: + if not isinstance(item, dict): + raise ValueError("history message is not an object") + pending = list(content) + while pending: + part = pending.pop() + if isinstance(part, list): + pending.extend(part) + elif isinstance(part, dict): + if part.get("type") == "image_media_ref": + parsed = ImageMediaRef( + media_id=part["media_id"], + mime_type=part["mime_type"], + width=part.get("width"), + height=part.get("height"), + byte_size=part["byte_size"], + detail=part.get("detail"), + version=part.get("version", 1), + image_id=part.get("image_id"), + ) + ImageMediaStore(media).read(parsed, {parsed.media_id}) + referenced.add(parsed.media_id) + else: + pending.extend(part.values()) + except Exception as exc: # noqa: BLE001 + raise RuntimeError("history parse failed; cleanup aborted") from exc + quarantine = Path( + args.quarantine or (media.parent / (media.name + ".quarantine")) + ) + if args.apply: + quarantine.mkdir(parents=True, exist_ok=True) + moved = 0 + for media_id in sorted( + {path.stem for path in media.iterdir() if path.suffix in {".bin", ".json"}} + ): + paths = [media / f"{media_id}.bin", media / f"{media_id}.json"] + if any(path.is_symlink() for path in paths if path.exists()): + continue + if media_id in referenced: + continue + existing_paths = [path for path in paths if path.exists()] + if len(existing_paths) != 2: + continue + try: + metadata = json.loads(paths[1].read_text()) + if ( + metadata.get("media_id") != media_id + or _digest(paths[0]) != media_id + ): + continue + except (OSError, json.JSONDecodeError): + continue + if args.apply: + targets = [quarantine / path.name for path in existing_paths] + if any(target.exists() for target in targets): + raise RuntimeError( + "quarantine target already exists; refusing overwrite" + ) + for path, target in zip(existing_paths, targets): + shutil.move(str(path), target) + moved += len(existing_paths) + if args.apply: + connection.commit() + else: + connection.rollback() + print( + json.dumps( + { + "dry_run": not args.apply, + "quarantined": moved, + "referenced": len(referenced), + } + ) + ) + finally: + connection.close() + + +def _restore(args: argparse.Namespace) -> None: + """Restore only when the database and media are unchanged since migration.""" + if not args.offline: + raise SystemExit("restore requires --offline") + backup = Path(args.restore).resolve() + db = Path(args.database).resolve() + backup_db = backup / db.name + manifest_path = backup / "manifest.json" + if not backup_db.is_file() or not manifest_path.is_file(): + raise SystemExit("backup is incomplete") + manifest = json.loads(manifest_path.read_text()) + if _digest(backup_db) != manifest.get("database"): + raise SystemExit("restore refused: archived database hash mismatch") + archived_media = backup / "media" + manifest_names: set[str] = set() + for item in manifest.get("media", []): + _manifest_media_name(item["name"]) + checksum = item.get("sha256", "") + if ( + item["name"] in manifest_names + or len(checksum) != 64 + or any(char not in "0123456789abcdef" for char in checksum) + ): + raise SystemExit("restore refused: invalid media manifest") + manifest_names.add(item["name"]) + source = archived_media / item["name"] + if ( + source.is_symlink() + or not source.is_file() + or _digest(source) != item["sha256"] + ): + raise SystemExit("restore refused: archived media hash mismatch") + current = _connect(db) + try: + if _logical_digest(current) != manifest.get("post_migration_logical_digest"): + raise SystemExit("restore refused: database changed since migration") + finally: + current.close() + media = Path(args.media).resolve() + for item in manifest.get("media", []): + target = media / item["name"] + if target.exists() and ( + target.is_symlink() or _digest(target) != item["sha256"] + ): + raise SystemExit(f"restore refused: conflicting media: {target.name}") + rollback = db.with_name(db.name + ".before-restore") + if rollback.exists(): + raise SystemExit("refusing to overwrite an existing pre-restore copy") + try: + staged_media: list[tuple[Path, Path]] = [] + staging = Path( + tempfile.mkdtemp(prefix="astrbot-restore-media-", dir=media.parent) + ) + for item in manifest.get("media", []): + source = archived_media / item["name"] + staged = staging / item["name"] + shutil.copy2(source, staged) + if _digest(staged) != item["sha256"]: + raise SystemExit("restore refused: staged media hash mismatch") + staged_media.append((staged, media / item["name"])) + media.mkdir(parents=True, exist_ok=True) + for staged, target in staged_media: + if not target.exists(): + staged.replace(target) + + rollback_connection = sqlite3.connect(rollback) + current = _connect(db) + try: + current.backup(rollback_connection) + rollback_connection.commit() + finally: + current.close() + rollback_connection.close() + destination = _connect(db) + archived = _connect(backup_db) + try: + archived.backup(destination) + destination.commit() + finally: + archived.close() + destination.close() + finally: + if "staging" in locals(): + shutil.rmtree(staging, ignore_errors=True) + print(json.dumps({"restored": str(db), "pre_restore_copy": str(rollback)})) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("database") + parser.add_argument("--media", required=True) + parser.add_argument("--conversation-id", action="append") + parser.add_argument("--apply", action="store_true") + parser.add_argument("--offline", action="store_true") + parser.add_argument("--backup") + parser.add_argument("--cleanup", action="store_true") + parser.add_argument("--quarantine") + parser.add_argument("--restore") + args = parser.parse_args() + if args.restore: + _restore(args) + elif args.cleanup: + _cleanup(args) + else: + _migrate(args) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_agent_runner_media_resolver.py b/tests/test_agent_runner_media_resolver.py index cd69f3aee1..4b5bbf289a 100644 --- a/tests/test_agent_runner_media_resolver.py +++ b/tests/test_agent_runner_media_resolver.py @@ -1,11 +1,20 @@ import base64 from io import BytesIO +from types import SimpleNamespace +from unittest.mock import AsyncMock import pytest from PIL import Image as PILImage +from astrbot.core.agent.runners.base import AgentState from astrbot.core.agent.runners.coze.coze_agent_runner import CozeAgentRunner +from astrbot.core.agent.runners.deerflow.deerflow_agent_runner import ( + DeerFlowAgentRunner, +) from astrbot.core.agent.runners.dify.dify_agent_runner import DifyAgentRunner +from astrbot.core.provider.entities import ProviderRequest +from astrbot.core.utils.image_media_store import ImageMediaStore +from astrbot.core.utils.media_utils import ImagePayloadTooLargeError def _png_data_url() -> tuple[str, bytes]: @@ -62,3 +71,108 @@ async def upload_file(self, file_data: bytes) -> str: assert file_id == "file-1" assert captured["file_data"] == image_bytes assert list(runner.file_id_cache["session-1"].values()) == ["file-1"] + + +@pytest.mark.asyncio +async def test_coze_history_reference_is_materialized_before_upload( + tmp_path, monkeypatch +): + import astrbot.core.agent.runners.coze.coze_agent_runner as coze_module + + image_ref, image_bytes = _png_data_url() + stored = ImageMediaStore(tmp_path / "media").put( + base64.b64decode(image_ref.split(",", 1)[1]) + ) + captured: dict[str, object] = {} + uploaded: list[str] = [] + upload_options = [] + + async def get_async(*_args, **_kwargs): + return "" + + monkeypatch.setattr(coze_module.sp, "get_async", get_async) + monkeypatch.setattr( + coze_module, + "get_astrbot_data_path", + lambda: str(tmp_path), + ) + + class _FakeCozeClient: + async def chat_messages(self, **kwargs): + captured.update(kwargs) + yield {"event": "conversation.message.completed", "data": {}} + yield {"event": "conversation.chat.completed", "data": {}} + + runner = CozeAgentRunner.__new__(CozeAgentRunner) + runner.req = ProviderRequest( + session_id="session-1", + contexts=[{"role": "user", "content": [stored.model_dump()]}], + ) + runner.auto_save_history = False + runner.api_client = _FakeCozeClient() + runner.bot_id = "bot-1" + runner.timeout = 10 + runner.streaming = False + runner._state = AgentState.RUNNING + + async def on_agent_done(*_args, **_kwargs): + return None + + runner.agent_hooks = SimpleNamespace( + on_agent_done=on_agent_done, + ) + runner.run_context = SimpleNamespace() + + async def capture_upload(image_url, _session_id, *, image_options=None): + uploaded.append(image_url) + upload_options.append(image_options) + return "file-1" + + runner._download_and_upload_image = capture_upload + + responses = [response async for response in runner._execute_coze_request()] + + assert len(responses) == 1 + assert len(uploaded) == 1 + assert upload_options[0].enabled is False + assert base64.b64decode(uploaded[0].split(",", 1)[1]) == image_bytes + assert captured["additional_messages"][0]["content"][0] == { + "type": "file", + "file_id": "file-1", + "file_url": uploaded[0], + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("runner_cls", "execute_method"), + [ + (DifyAgentRunner, "_execute_dify_request"), + (CozeAgentRunner, "_execute_coze_request"), + (DeerFlowAgentRunner, "_execute_deerflow_request"), + ], +) +async def test_resource_failure_is_not_converted_to_agent_response( + runner_cls, execute_method +): + failure = ImagePayloadTooLargeError("image exceeds the request budget") + + async def failed_request(): + raise failure + yield # pragma: no cover + + runner = runner_cls.__new__(runner_cls) + runner.req = ProviderRequest(prompt="hello") + runner._state = AgentState.IDLE + runner.agent_hooks = SimpleNamespace(on_agent_begin=AsyncMock()) + runner.run_context = SimpleNamespace() + setattr(runner, execute_method, failed_request) + if runner_cls in {DifyAgentRunner, CozeAgentRunner}: + runner.api_client = SimpleNamespace(close=AsyncMock()) + + with pytest.raises(ImagePayloadTooLargeError) as caught: + async for _ in runner.step(): + pass + + assert caught.value is failure + assert runner._state == AgentState.ERROR diff --git a/tests/test_computer_fs_tools.py b/tests/test_computer_fs_tools.py index a2f26670d8..7e224e1bab 100644 --- a/tests/test_computer_fs_tools.py +++ b/tests/test_computer_fs_tools.py @@ -184,12 +184,6 @@ def _setup_local_fs_tools( "get_astrbot_temp_path", lambda: str(temp_root), ) - monkeypatch.setattr( - file_read_utils, - "get_astrbot_temp_path", - lambda: str(temp_root), - ) - booter = LocalBooter() async def _fake_get_booter(_ctx, _umo): @@ -861,8 +855,47 @@ async def test_file_read_tool_returns_image_call_tool_result_for_images( assert isinstance(result, CallToolResult) assert len(result.content) == 1 assert isinstance(result.content[0], ImageContent) - assert result.content[0].mimeType == "image/jpeg" - assert base64.b64decode(result.content[0].data).startswith(b"\xff\xd8\xff") + assert result.content[0].mimeType == "image/png" + assert base64.b64decode(result.content[0].data) == image_path.read_bytes() + + +@pytest.mark.asyncio +async def test_local_file_read_defers_image_preparation_to_tool_loop( + monkeypatch: pytest.MonkeyPatch, + tmp_path, +): + workspace = _setup_local_fs_tools(monkeypatch, tmp_path) + image_path = workspace / "sample.png" + Image.new("RGB", (32, 16), color=(255, 0, 0)).save(image_path, format="PNG") + result = await fs_tools.FileReadTool().call( + _make_context(), + path="sample.png", + ) + + assert isinstance(result, CallToolResult) + assert result.content[0].data + assert base64.b64decode(result.content[0].data) == image_path.read_bytes() + + +@pytest.mark.asyncio +async def test_file_read_propagates_image_memory_error( + monkeypatch: pytest.MonkeyPatch, + tmp_path, +): + workspace = _setup_local_fs_tools(monkeypatch, tmp_path) + image_path = workspace / "sample.png" + Image.new("RGB", (32, 16), color=(255, 0, 0)).save(image_path, format="PNG") + + async def fail_read(*_args, **_kwargs): + raise MemoryError("image preparation exhausted memory") + + monkeypatch.setattr(file_read_utils, "_read_local_file_bytes", fail_read) + + with pytest.raises(MemoryError, match="exhausted memory"): + await fs_tools.FileReadTool().call( + _make_context(), + path="sample.png", + ) @pytest.mark.asyncio diff --git a/tests/test_fastapi_v1_dashboard.py b/tests/test_fastapi_v1_dashboard.py index f9262204ff..f8d4bad1e3 100644 --- a/tests/test_fastapi_v1_dashboard.py +++ b/tests/test_fastapi_v1_dashboard.py @@ -1671,6 +1671,46 @@ async def test_v1_conversation_detail_requires_user_id( assert response.status_code == 422 +@pytest.mark.asyncio +async def test_v1_conversation_media_route_requires_auth_and_owner( + asgi_client: httpx.AsyncClient, + monkeypatch, +): + from astrbot.dashboard.services.conversation_service import ( + ConversationMedia, + ConversationService, + ConversationServiceError, + ) + + async def preview(self, user_id: str, cid: str, media_id: str): + if user_id != "owner" or cid != "cid" or media_id != "a" * 64: + raise ConversationServiceError("对话不存在") + return ConversationMedia(b"valid-image", "image/png") + + monkeypatch.setattr(ConversationService, "get_conversation_media", preview) + path = "/api/v1/conversations/cid/media/" + "a" * 64 + assert (await asgi_client.get(path, params={"user_id": "owner"})).status_code == 401 + + response = await asgi_client.get( + path, params={"user_id": "owner"}, headers=_jwt_headers() + ) + assert response.status_code == 200 + assert response.headers["content-type"] == "image/png" + assert response.content == b"valid-image" + + other_owner = await asgi_client.get( + path, params={"user_id": "other"}, headers=_jwt_headers() + ) + assert other_owner.status_code == 400 + + unreferenced = await asgi_client.get( + "/api/v1/conversations/cid/media/" + "b" * 64, + params={"user_id": "owner"}, + headers=_jwt_headers(), + ) + assert unreferenced.status_code == 400 + + @pytest.mark.asyncio async def test_dashboard_alias_conversation_detail_uses_fastapi_service( asgi_client: httpx.AsyncClient, diff --git a/tests/test_media_utils.py b/tests/test_media_utils.py index 46dd1f66b4..08402cfdde 100644 --- a/tests/test_media_utils.py +++ b/tests/test_media_utils.py @@ -324,6 +324,30 @@ async def test_compress_image_preserves_alpha_png(tmp_path, monkeypatch): compressed_path.unlink(missing_ok=True) +@pytest.mark.asyncio +async def test_compress_image_rejects_oversized_encoded_payload(tmp_path, monkeypatch): + from PIL import Image as PILImage + + temp_dir = tmp_path / "temp" + monkeypatch.setattr(media_utils, "get_astrbot_temp_path", lambda: str(temp_dir)) + image_path = tmp_path / "high_entropy.png" + image = PILImage.new("RGBA", (32, 32)) + image.putdata( + [ + (index % 256, (index * 7) % 256, (index * 13) % 256, 255) + for index in range(1024) + ] + ) + image.save(image_path, format="PNG") + + with pytest.raises(media_utils.ImagePayloadTooLargeError, match="encoding limit"): + await media_utils.compress_image( + str(image_path), max_size=128, max_encoded_bytes=1 + ) + + assert not temp_dir.exists() or not list(temp_dir.iterdir()) + + @pytest.mark.asyncio async def test_compress_image_keeps_animated_gif(tmp_path, monkeypatch): from PIL import Image as PILImage @@ -343,7 +367,7 @@ async def test_compress_image_keeps_animated_gif(tmp_path, monkeypatch): compressed_path = await media_utils.compress_image(str(image_path), max_size=2) assert compressed_path == str(image_path) - assert not list(temp_dir.iterdir()) + assert not temp_dir.exists() or not list(temp_dir.iterdir()) @pytest.mark.asyncio diff --git a/tests/test_openai_responses_source.py b/tests/test_openai_responses_source.py index 6b2d5e6718..04192c4041 100644 --- a/tests/test_openai_responses_source.py +++ b/tests/test_openai_responses_source.py @@ -1,13 +1,16 @@ +import io import json from types import SimpleNamespace import pytest from openai.types.responses import Response +from PIL import Image from astrbot.core.config.default import CONFIG_METADATA_2 from astrbot.core.provider.sources.openai_responses_source import ( ProviderOpenAIResponses, ) +from astrbot.core.utils.image_media_store import ImageMediaStore def _make_provider(overrides: dict | None = None) -> ProviderOpenAIResponses: @@ -208,6 +211,32 @@ async def test_prepare_payload_replays_full_history_without_server_state(): assert "conversation" not in payloads +@pytest.mark.asyncio +async def test_prepare_payload_materializes_durable_image_reference( + tmp_path, monkeypatch +): + output = io.BytesIO() + Image.new("RGB", (3, 2), "red").save(output, format="PNG") + data = output.getvalue() + store = ImageMediaStore(tmp_path / "media") + ref = store.put(data, "image/png", detail="high") + monkeypatch.setattr( + "astrbot.core.provider.sources.openai_responses_source.get_astrbot_data_path", + lambda: str(tmp_path), + ) + provider = _make_provider() + + payloads, context = await provider._prepare_chat_payload( + prompt=None, + contexts=[{"role": "user", "content": [ref.model_dump()]}], + ) + + image = context[0]["content"][0]["image_url"] + assert image["url"].startswith("data:image/png;base64,") + assert image["detail"] == "high" + assert payloads["input"][0]["content"][0]["type"] == "input_image" + + @pytest.mark.asyncio async def test_query_flattens_tools_and_enforces_stateless_body(monkeypatch): provider = _make_provider( diff --git a/tests/test_openai_source.py b/tests/test_openai_source.py index 76a65171be..1fe68e317c 100644 --- a/tests/test_openai_source.py +++ b/tests/test_openai_source.py @@ -1,6 +1,7 @@ import base64 import builtins from io import BytesIO +from pathlib import Path from types import SimpleNamespace import httpx @@ -16,8 +17,6 @@ from astrbot.core.provider.entities import LLMResponse from astrbot.core.provider.sources.groq_source import ProviderGroq from astrbot.core.provider.sources.openai_source import ProviderOpenAIOfficial -from pathlib import Path - from astrbot.core.utils.media_utils import ResolvedMediaData, file_uri_to_path @@ -774,10 +773,12 @@ async def fake_resolve_media_ref_to_base64_data( *, media_type: str, strict: bool = False, + image_options=None, ) -> ResolvedMediaData: assert media_ref == "https://example.com/quoted.png" assert media_type == "image" assert strict is False + assert image_options is not None return ResolvedMediaData(base64_data="abcd", mime_type="image/png") monkeypatch.setattr( @@ -1044,9 +1045,15 @@ async def test_materialize_context_image_parts_returns_new_messages(monkeypatch) {"role": "assistant", "content": "plain text"}, ] - async def fake_resolve(image_url: str, *, image_detail: str | None = None): + async def fake_resolve( + image_url: str, + *, + image_detail: str | None = None, + options=None, + ): assert image_url == "https://example.com/quoted.png" assert image_detail == "high" + assert options is not None return { "type": "image_url", "image_url": { @@ -1077,6 +1084,35 @@ async def fake_resolve(image_url: str, *, image_detail: str | None = None): await provider.terminate() +@pytest.mark.asyncio +async def test_materialize_context_keeps_existing_data_image_without_reencoding( + monkeypatch, +): + provider = _make_provider() + try: + image_part = { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,already-prepared", + "detail": "high", + "id": "request-image", + }, + } + + async def fail_if_resolved(*_args, **_kwargs): + raise AssertionError("an existing data URL must not be re-encoded") + + monkeypatch.setattr(provider, "_resolve_image_part", fail_if_resolved) + + materialized = await provider._materialize_context_image_parts( + [{"role": "user", "content": [image_part]}] + ) + + assert materialized[0]["content"][0] == image_part + finally: + await provider.terminate() + + @pytest.mark.asyncio async def test_encode_image_bs64_missing_file_raises(tmp_path): provider = _make_provider() diff --git a/tests/test_process_stage_images.py b/tests/test_process_stage_images.py index 44c2aba77b..7c9f60aef6 100644 --- a/tests/test_process_stage_images.py +++ b/tests/test_process_stage_images.py @@ -19,6 +19,7 @@ TextPart, dump_messages_with_checkpoints, ) +from astrbot.core.agent.runners import tool_loop_agent_runner from astrbot.core.config.default import DEFAULT_CONFIG from astrbot.core.message.components import Image, Plain, Reply from astrbot.core.pipeline.preprocess_stage import stage as preprocess @@ -35,6 +36,10 @@ from astrbot.core.provider.provider import Provider from astrbot.core.star.star_handler import EventType from astrbot.core.utils import media_utils as media +from astrbot.core.utils.image_media_store import ( + ImageMediaStore, + materialize_image_media_refs, +) def make_event(parts=None, text="hello", session="images"): @@ -366,6 +371,11 @@ async def test_profile_reload_and_concurrent_requests(harness, tmp_path): async def test_plugin_request_extra_metadata_hook_and_history( harness, tmp_path, monkeypatch ): + monkeypatch.setattr( + tool_loop_agent_runner, + "get_astrbot_data_path", + lambda: str(tmp_path / "data"), + ) source = source_image(tmp_path) replacement = source_image(tmp_path, "BMP") extra = ImageURLPart( @@ -438,7 +448,14 @@ async def hook(event, kind, *args): == historical[0]["content"][0]["image_url"]["url"] ) assert req.contexts == historical - images = [p for p in saved[-1]["content"] if p["type"] == "image_url"] + refs = [p for p in saved[-1]["content"] if p["type"] == "image_media_ref"] + assert refs and all(len(p["media_id"]) == 64 for p in refs) + materialized = await materialize_image_media_refs( + [saved[-1]], ImageMediaStore(tmp_path / "data" / "media") + ) + images = [ + p for p in materialized[0]["content"] if p["type"] == "image_url" + ] assert images and all( p["image_url"]["url"].startswith("data:image/jpeg;base64,") for p in images ) @@ -455,7 +472,7 @@ async def hook(event, kind, *args): from astrbot.core.provider.sources.anthropic_source import ProviderAnthropic anthropic = object.__new__(ProviderAnthropic) - _, payload = anthropic._prepare_payload([saved[-1]]) + _, payload = anthropic._prepare_payload(materialized) visual = [part for part in payload[0]["content"] if part["type"] == "image"] assert len(visual) == len(images) assert all(part["source"]["media_type"] == "image/jpeg" for part in visual) diff --git a/tests/test_tool_loop_agent_runner.py b/tests/test_tool_loop_agent_runner.py index 180e0edf2d..bd68e7ca22 100644 --- a/tests/test_tool_loop_agent_runner.py +++ b/tests/test_tool_loop_agent_runner.py @@ -133,7 +133,10 @@ async def generator(): content=[ ImageContent( type="image", - data="dGVzdA==", + data=( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk" + "+A8AAQUBAScY42YAAAAASUVORK5CYII=" + ), mimeType="image/png", ), TextContent(type="text", text="直播间标题:新游首发:零~红蝶~"), @@ -938,10 +941,26 @@ def fake_save_image( } ) return SimpleNamespace( - file_path=f"/tmp/{tool_call_id}_{index}.png", mime_type=mime_type + file_path=f"/tmp/{tool_call_id}_{index}.png", + mime_type=mime_type, + tool_name=tool_name, ) monkeypatch.setattr(tool_image_cache, "save_image", fake_save_image) + from astrbot.core.utils.media_utils import ResolvedMediaData + + monkeypatch.setattr( + "astrbot.core.agent.runners.tool_loop_agent_runner.prepare_image_source", + AsyncMock( + return_value=ResolvedMediaData( + base64_data=( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk" + "+A8AAQUBAScY42YAAAAASUVORK5CYII=" + ), + mime_type="image/png", + ) + ), + ) await runner.reset( provider=mock_provider, @@ -965,7 +984,10 @@ def fake_save_image( assert "直播间标题:新游首发:零~红蝶~" in content assert saved_images == [ { - "base64_data": "dGVzdA==", + "base64_data": ( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk" + "+A8AAQUBAScY42YAAAAASUVORK5CYII=" + ), "tool_call_id": "call_123", "tool_name": "test_tool", "index": 0, @@ -1252,6 +1274,297 @@ async def test_same_tool_streak_resets_after_switching_tools( assert level_2_notice in content +@pytest.mark.asyncio +@pytest.mark.parametrize("error_type", [MemoryError]) +async def test_resource_exhaustion_does_not_try_fallback( + runner, provider_request, mock_tool_executor, mock_hooks, error_type +): + primary = MockProvider() + fallback = MockProvider() + primary.text_chat = AsyncMock(side_effect=error_type("resource exhausted")) + await runner.reset( + provider=primary, + request=provider_request, + run_context=ContextWrapper(context=None), + tool_executor=mock_tool_executor, + agent_hooks=mock_hooks, + streaming=False, + fallback_providers=[fallback], + ) + with pytest.raises(error_type): + async for _ in runner._iter_llm_responses_with_fallback(): + pass + assert fallback.call_count == 0 + + +@pytest.mark.asyncio +async def test_http_413_is_readable_and_not_retried( + runner, provider_request, mock_tool_executor, mock_hooks +): + import httpx + + from astrbot.core.exceptions import ProviderRequestTooLargeError + + primary = MockProvider() + fallback = MockProvider() + request = httpx.Request("POST", "https://invalid.test/chat") + response = httpx.Response(413, request=request) + error = httpx.HTTPStatusError("too large", request=request, response=response) + primary.text_chat = AsyncMock(side_effect=error) + await runner.reset( + provider=primary, + request=provider_request, + run_context=ContextWrapper(context=None), + tool_executor=mock_tool_executor, + agent_hooks=mock_hooks, + fallback_providers=[fallback], + ) + with pytest.raises(ProviderRequestTooLargeError, match="HTTP 413"): + async for _ in runner._iter_llm_responses_with_fallback(): + pass + assert primary.text_chat.await_count == 1 + assert fallback.call_count == 0 + + +@pytest.mark.asyncio +async def test_tool_memory_failure_is_not_returned_as_empty_tool_error( + runner, provider_request, mock_hooks +): + failure = MemoryError() + + class FailedExecutor: + @classmethod + def execute(cls, *args, **kwargs): + async def results(): + raise failure + yield # pragma: no cover + + return results() + + provider = MockProvider() + await runner.reset( + provider=provider, + request=provider_request, + run_context=ContextWrapper(context=None), + tool_executor=FailedExecutor, + agent_hooks=mock_hooks, + ) + with pytest.raises(MemoryError) as caught: + async for _ in runner.step_until_done(3): + pass + assert caught.value is failure + assert provider.call_count == 1 + + +@pytest.mark.asyncio +async def test_real_tool_image_uses_preparation_and_reference_storage( + runner, provider_request, mock_hooks, tmp_path, monkeypatch +): + import base64 + import io + + from mcp.types import CallToolResult, ImageContent + from PIL import Image + + import astrbot.core.agent.runners.tool_loop_agent_runner as runner_module + from astrbot.core.agent.message import ImageMediaRefPart + from astrbot.core.agent.tool_image_cache import tool_image_cache + from astrbot.core.utils.image_media_store import ImageMediaStore + + output = io.BytesIO() + with Image.new("RGBA", (24, 16), (60, 20, 30, 128)) as image: + image.save(output, "WEBP", lossless=True) + encoded = base64.b64encode(output.getvalue()).decode() + + class ImageExecutor: + @classmethod + def execute(cls, *args, **kwargs): + async def results(): + yield CallToolResult( + content=[ + ImageContent(type="image", data=encoded, mimeType="image/webp") + ] + ) + + return results() + + monkeypatch.setattr(tool_image_cache, "_cache_dir", str(tmp_path / "tool-cache")) + original_prepare = runner_module.prepare_image_source + captured_options = [] + + async def capture_prepare(image_ref, **kwargs): + captured_options.append(kwargs["options"]) + return await original_prepare(image_ref, **kwargs) + + monkeypatch.setattr(runner_module, "prepare_image_source", capture_prepare) + provider = MockProvider() + provider.provider_settings = { + "image_compress_options": { + "max_size": 512, + "quality": 42, + "max_encoded_bytes": 2 * 1024 * 1024, + } + } + provider.max_calls_before_normal_response = 1 + await runner.reset( + provider=provider, + request=provider_request, + run_context=ContextWrapper(context=None), + tool_executor=ImageExecutor, + agent_hooks=mock_hooks, + image_media_store=ImageMediaStore(tmp_path / "media"), + ) + async for _ in runner.step_until_done(3): + pass + images = [ + part + for message in runner.run_context.messages + if isinstance(message.content, list) + for part in message.content + if isinstance(part, ImageMediaRefPart) + ] + assert len(images) == 1 + assert images[0].mime_type == "image/webp" + assert images[0].image_id.endswith("call_123_0.webp") + assert captured_options[0].max_size == 512 + assert captured_options[0].quality == 42 + assert captured_options[0].max_encoded_bytes == 2 * 1024 * 1024 + assert provider.call_count == 2 + + +@pytest.mark.asyncio +async def test_new_media_refs_do_not_rewrite_existing_inline_history( + runner, mock_tool_executor, mock_hooks, tmp_path +): + import base64 + import io + + from PIL import Image + + from astrbot.core.agent.message import ImageMediaRefPart + from astrbot.core.utils.image_media_store import ImageMediaStore + + output = io.BytesIO() + with Image.new("RGB", (10, 8), (10, 20, 30)) as image: + image.save(output, "PNG") + uri = "data:image/png;base64," + base64.b64encode(output.getvalue()).decode() + old_history = [ + {"role": "user", "content": [{"type": "image_url", "image_url": {"url": uri}}]}, + {"role": "assistant", "content": "Old image"}, + ] + request = ProviderRequest( + prompt="new image", image_urls=[uri], contexts=old_history + ) + provider = MockProvider() + provider.text_chat = AsyncMock( + return_value=LLMResponse(role="assistant", completion_text="ok") + ) + await runner.reset( + provider=provider, + request=request, + run_context=ContextWrapper(context=None), + tool_executor=mock_tool_executor, + agent_hooks=mock_hooks, + image_media_store=ImageMediaStore(tmp_path / "media"), + ) + assert runner.run_context.messages[0].content[0].image_url.url == uri + assert isinstance(runner.run_context.messages[-1].content[-1], ImageMediaRefPart) + async for _ in runner._iter_llm_responses(): + pass + sent = provider.text_chat.await_args.kwargs["contexts"] + assert sent[0].content[0].image_url.url == uri + assert sent[-1].content[-1].image_url.url == uri + assert isinstance(runner.run_context.messages[-1].content[-1], ImageMediaRefPart) + assert old_history[0]["content"][0]["image_url"]["url"] == uri + + +@pytest.mark.asyncio +async def test_runner_uses_default_durable_media_store_for_new_images( + runner, mock_tool_executor, mock_hooks, tmp_path, monkeypatch +): + import base64 + import io + + from PIL import Image + + from astrbot.core.agent.message import ImageMediaRefPart + + output = io.BytesIO() + with Image.new("RGB", (10, 8), (10, 20, 30)) as image: + image.save(output, "PNG") + uri = "data:image/png;base64," + base64.b64encode(output.getvalue()).decode() + monkeypatch.setattr( + "astrbot.core.agent.runners.tool_loop_agent_runner.get_astrbot_data_path", + lambda: str(tmp_path), + ) + request = ProviderRequest(prompt="new image", image_urls=[uri]) + + await runner.reset( + provider=MockProvider(), + request=request, + run_context=ContextWrapper(context=None), + tool_executor=mock_tool_executor, + agent_hooks=mock_hooks, + ) + + image_part = runner.run_context.messages[-1].content[-1] + assert isinstance(image_part, ImageMediaRefPart) + assert image_part.image_id is None + assert (tmp_path / "media" / f"{image_part.media_id}.bin").exists() + + +@pytest.mark.asyncio +async def test_context_selection_does_not_open_out_of_window_images( + runner, mock_tool_executor, mock_hooks, tmp_path, monkeypatch +): + from PIL import Image + + from astrbot.core.utils.image_media_store import ImageMediaStore + + image_path = tmp_path / "image.png" + Image.new("RGB", (8, 8), "green").save(image_path) + store = ImageMediaStore(tmp_path / "media") + ref = store.put(image_path.read_bytes()) + missing = {**ref.model_dump(), "media_id": "f" * 64, "byte_size": 100_000_000} + request = ProviderRequest( + prompt="current question", + contexts=[ + {"role": "user", "content": [missing]}, + {"role": "assistant", "content": "old answer"}, + {"role": "user", "content": "another old question"}, + {"role": "assistant", "content": "another old answer"}, + {"role": "user", "content": [ref.model_dump()]}, + {"role": "assistant", "content": "recent answer"}, + ], + ) + provider = MockProvider() + provider.text_chat = AsyncMock( + return_value=LLMResponse(role="assistant", completion_text="ok") + ) + opened = [] + original_read = store.read + + def read(reference, allowed_ids): + opened.append(reference.media_id) + return original_read(reference, allowed_ids) + + monkeypatch.setattr(store, "read", read) + await runner.reset( + provider=provider, + request=request, + run_context=ContextWrapper(context=None), + tool_executor=mock_tool_executor, + agent_hooks=mock_hooks, + image_media_store=store, + enforce_max_turns=2, + ) + async for _ in runner.step(): + pass + assert opened == [ref.media_id] + assert provider.text_chat.await_count == 1 + assert request.contexts[0]["content"][0] == missing + + @pytest.mark.asyncio async def test_fallback_provider_used_when_primary_raises( runner, provider_request, mock_tool_executor, mock_hooks @@ -1656,9 +1969,19 @@ async def test_follow_up_ticket_not_consumed_when_no_next_tool_call( @pytest.mark.asyncio @pytest.mark.parametrize("streaming", [False, True]) -async def test_skills_like_requery_passes_extra_user_content_parts(streaming): +async def test_skills_like_requery_passes_extra_user_content_parts(streaming, tmp_path): """skills-like 模式 re-query 时应传递 extra_user_content_parts(如 image_caption)""" + from PIL import Image + from astrbot.core.agent.message import TextPart + from astrbot.core.utils.image_media_store import ImageMediaStore + + image_path = tmp_path / "image.png" + Image.new("RGB", (8, 8), "red").save(image_path) + image_bytes = image_path.read_bytes() + store = ImageMediaStore(tmp_path / "media") + ref = store.put(image_bytes, detail="high", image_id="requery-image") + historical_image = {"role": "user", "content": [ref.model_dump()]} captured_kwargs = {} @@ -1706,7 +2029,7 @@ async def text_chat(self, **kwargs) -> LLMResponse: req = ProviderRequest( prompt="看看这张图", func_tool=tool_set, - contexts=[], + contexts=[historical_image, {"role": "assistant", "content": "ack"}], extra_user_content_parts=[caption_part], ) @@ -1723,6 +2046,7 @@ async def text_chat(self, **kwargs) -> LLMResponse: agent_hooks=MockHooks(), tool_schema_mode="skills_like", streaming=streaming, + image_media_store=store, ) async for _ in runner.step(): @@ -1735,6 +2059,13 @@ async def text_chat(self, **kwargs) -> LLMResponse: parts = captured_kwargs["extra_user_content_parts"] assert len(parts) == 1 assert parts[0].text == "一张猫的照片" + import base64 + + sent_image = captured_kwargs["contexts"][0]["content"][0]["image_url"] + assert base64.b64decode(sent_image["url"].split(",", 1)[1]) == image_bytes + assert sent_image["detail"] == "high" + assert sent_image["id"] == "requery-image" + assert historical_image["content"][0] == ref.model_dump() @pytest.mark.asyncio diff --git a/tests/unit/test_astr_main_agent.py b/tests/unit/test_astr_main_agent.py index 3319e99f3e..88f740d5a9 100644 --- a/tests/unit/test_astr_main_agent.py +++ b/tests/unit/test_astr_main_agent.py @@ -20,6 +20,7 @@ from astrbot.core.config.agent_runner import resolve_context_compression_config from astrbot.core.conversation_mgr import Conversation from astrbot.core.cron.manager import CronJobManager +from astrbot.core.exceptions import ProviderRequestTooLargeError from astrbot.core.message.components import File, Image, Plain, Reply, Video from astrbot.core.platform.astr_message_event import AstrMessageEvent from astrbot.core.platform.platform_metadata import PlatformMetadata @@ -30,6 +31,7 @@ from astrbot.core.skills.skill_manager import SkillInfo from astrbot.core.star.context import Context from astrbot.core.star.star import StarMetadata +from astrbot.core.utils.media_utils import ImagePayloadTooLargeError @pytest.fixture @@ -44,6 +46,27 @@ def mock_provider(): return provider +@pytest.mark.asyncio +@pytest.mark.parametrize( + "error", + [ + MemoryError("oom"), + ImagePayloadTooLargeError("too large"), + ProviderRequestTooLargeError("413"), + ], +) +async def test_image_caption_propagates_resource_errors(monkeypatch, error): + req = ProviderRequest(image_urls=["image.png"]) + + async def fail(*_args, **_kwargs): + raise error + + monkeypatch.setattr(ama, "_request_img_caption", fail) + with pytest.raises(type(error)): + await ama._ensure_img_caption(None, req, {}, MagicMock(), "caption") + assert req.image_urls == [] + + @pytest.fixture def mock_context(): """Create a mock Context.""" diff --git a/tests/unit/test_context_image_budget.py b/tests/unit/test_context_image_budget.py new file mode 100644 index 0000000000..809001b51f --- /dev/null +++ b/tests/unit/test_context_image_budget.py @@ -0,0 +1,153 @@ +"""Image-byte budgets must not change history or token accounting.""" + +import base64 +import copy +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +from PIL import Image + +from astrbot.core.agent.context.compressor import LLMSummaryCompressor +from astrbot.core.agent.context.image_budget import ( + get_image_encoded_byte_limit, + validate_context_image_bytes, +) +from astrbot.core.agent.message import Message +from astrbot.core.exceptions import ProviderRequestTooLargeError +from astrbot.core.utils.image_media_store import ImageMediaStore +from astrbot.core.utils.media_utils import ImagePayloadTooLargeError + + +@pytest.mark.parametrize("limit", [None, True, False, 0, -1, "100"]) +def test_invalid_provider_image_limit_uses_default(limit): + assert ( + get_image_encoded_byte_limit( + {"image_compress_options": {"max_encoded_bytes": limit}} + ) + == 4 * 1024 * 1024 + ) + + +def test_provider_image_limit_is_respected(): + assert ( + get_image_encoded_byte_limit( + {"image_compress_options": {"max_encoded_bytes": 16 * 1024 * 1024}} + ) + == 16 * 1024 * 1024 + ) + + +@pytest.mark.parametrize("as_models", [False, True]) +def test_byte_budget_does_not_rewrite_history(as_models): + history = [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,AAAA"}, + } + ], + } + ] + messages = ( + [Message.model_validate(item) for item in history] if as_models else history + ) + before = copy.deepcopy(messages) + assert validate_context_image_bytes(messages, 4) == 4 + with pytest.raises(ImagePayloadTooLargeError): + validate_context_image_bytes(messages, 3) + assert messages == before + + +def test_reference_size_is_checked_without_opening_media(): + history = [ + { + "role": "user", + "content": [ + { + "type": "image_media_ref", + "media_id": "a" * 64, + "byte_size": 12, + "mime_type": "image/png", + "version": 1, + } + ], + } + ] + assert validate_context_image_bytes(history, 16) == 16 + with pytest.raises(ImagePayloadTooLargeError): + validate_context_image_bytes(history, 15) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("vision", [True, False]) +async def test_summary_loads_only_its_selected_images(tmp_path, monkeypatch, vision): + image_path = tmp_path / "old.png" + Image.new("RGB", (8, 8), "blue").save(image_path) + original = image_path.read_bytes() + store = ImageMediaStore(tmp_path / "media") + old_ref = store.put(original, detail="high", image_id="old-image") + messages = [ + Message.model_validate({"role": "user", "content": [old_ref.model_dump()]}), + Message(role="assistant", content="old answer"), + Message.model_validate( + { + "role": "user", + "content": [{**old_ref.model_dump(), "media_id": "f" * 64}], + } + ), + ] + before = [message.model_dump() for message in messages] + opened = [] + original_read = store.read + + def read(ref, allowed_ids): + opened.append(ref.media_id) + return original_read(ref, allowed_ids) + + monkeypatch.setattr(store, "read", read) + provider = SimpleNamespace( + provider_config={"modalities": ["text", "image"] if vision else ["text"]}, + provider_settings={}, + text_chat=AsyncMock(return_value=SimpleNamespace(completion_text="summary")), + ) + result = await LLMSummaryCompressor( + provider, keep_recent_ratio=0, image_media_store=store + )(messages) + assert opened == ([old_ref.media_id] if vision else []) + assert result[-1] is messages[-1] + assert [message.model_dump() for message in messages] == before + if vision: + sent = provider.text_chat.call_args.kwargs["contexts"][0] + image = sent["content"][0]["image_url"] + assert base64.b64decode(image["url"].split(",", 1)[1]) == original + assert image["detail"] == "high" + assert image["id"] == "old-image" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "error", + [ + MemoryError(), + ImagePayloadTooLargeError("oversize"), + ProviderRequestTooLargeError("413"), + ], +) +async def test_summary_resource_failure_is_not_swallowed(error): + provider = SimpleNamespace( + provider_config={"modalities": ["text"]}, + provider_settings={}, + text_chat=AsyncMock(side_effect=error), + ) + messages = [ + Message(role="user", content="old question"), + Message(role="assistant", content="old answer"), + Message(role="user", content="current question"), + ] + with pytest.raises(type(error)) as caught: + await LLMSummaryCompressor(provider, keep_recent_ratio=0)(messages) + assert caught.value is error + assert provider.text_chat.await_count == 1 diff --git a/tests/unit/test_conversation_media_api.py b/tests/unit/test_conversation_media_api.py new file mode 100644 index 0000000000..e4b0ac96fe --- /dev/null +++ b/tests/unit/test_conversation_media_api.py @@ -0,0 +1,149 @@ +import io +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest +from PIL import Image + +from astrbot.dashboard.services.conversation_service import ( + ConversationService, + ConversationServiceError, +) + + +def _image_bytes() -> bytes: + output = io.BytesIO() + Image.new("RGB", (2, 2), "red").save(output, format="PNG") + return output.getvalue() + + +@pytest.mark.asyncio +async def test_media_preview_requires_database_owner_and_reference( + tmp_path, monkeypatch +): + monkeypatch.setattr( + "astrbot.dashboard.services.conversation_service.get_astrbot_data_path", + lambda: str(tmp_path), + ) + data = _image_bytes() + from astrbot.core.utils.image_media_store import ImageMediaStore + + ref = ImageMediaStore(Path(tmp_path) / "media").put(data, "image/png") + history = json.dumps([{"content": [ref.model_dump()]}]) + conversation = SimpleNamespace(user_id="owner", history=history) + + class Db: + async def get_conversation_by_id(self, cid): + return conversation if cid == "cid" else None + + service = ConversationService(Db(), SimpleNamespace(conversation_manager=None)) + result = await service.get_conversation_media("owner", "cid", ref.media_id) + assert result.data == data + + with pytest.raises(ConversationServiceError): + await service.get_conversation_media("other", "cid", ref.media_id) + with pytest.raises(ConversationServiceError): + await service.get_conversation_media("owner", "cid", "0" * 64) + + +@pytest.mark.asyncio +async def test_media_preview_does_not_cross_conversation_reference_boundary( + tmp_path, monkeypatch +): + monkeypatch.setattr( + "astrbot.dashboard.services.conversation_service.get_astrbot_data_path", + lambda: str(tmp_path), + ) + from astrbot.core.utils.image_media_store import ImageMediaStore + + ref = ImageMediaStore(Path(tmp_path) / "media").put(_image_bytes(), "image/png") + conversation = SimpleNamespace(user_id="owner", history=json.dumps([])) + + class Db: + async def get_conversation_by_id(self, cid): + return conversation if cid == "different-conversation" else None + + service = ConversationService(Db(), SimpleNamespace(conversation_manager=None)) + with pytest.raises(ConversationServiceError, match="媒体不存在"): + await service.get_conversation_media( + "owner", "different-conversation", ref.media_id + ) + + +@pytest.mark.asyncio +async def test_media_preview_missing_or_corrupt_does_not_expose_path( + tmp_path, monkeypatch +): + monkeypatch.setattr( + "astrbot.dashboard.services.conversation_service.get_astrbot_data_path", + lambda: str(tmp_path), + ) + from astrbot.core.utils.image_media_store import ImageMediaStore + + store = ImageMediaStore(Path(tmp_path) / "media") + ref = store.put(_image_bytes(), "image/png") + conversation = SimpleNamespace( + user_id="owner", history=json.dumps([{"content": [ref.model_dump()]}]) + ) + + class Db: + async def get_conversation_by_id(self, cid): + return conversation + + service = ConversationService(Db(), SimpleNamespace(conversation_manager=None)) + (Path(tmp_path) / "media" / f"{ref.media_id}.bin").write_bytes(b"corrupt") + with pytest.raises(ConversationServiceError) as exc_info: + await service.get_conversation_media("owner", "cid", ref.media_id) + assert str(tmp_path) not in str(exc_info.value) + + (Path(tmp_path) / "media" / f"{ref.media_id}.bin").unlink() + with pytest.raises(ConversationServiceError) as exc_info: + await service.get_conversation_media("owner", "cid", ref.media_id) + assert str(tmp_path) not in str(exc_info.value) + + +@pytest.mark.asyncio +async def test_export_materializes_refs_with_detail_and_image_id(tmp_path, monkeypatch): + monkeypatch.setattr( + "astrbot.dashboard.services.conversation_service.get_astrbot_data_path", + lambda: str(tmp_path), + ) + from astrbot.core.utils.image_media_store import ImageMediaStore + + store = ImageMediaStore(Path(tmp_path) / "media") + ref = store.put(_image_bytes(), "image/png", "high") + ref = type(ref)( + ref.media_id, + ref.mime_type, + ref.width, + ref.height, + ref.byte_size, + "high", + 1, + "img-1", + ) + conversation = SimpleNamespace( + user_id="owner", + history=json.dumps([{"role": "user", "content": [ref.model_dump()]}]), + platform_id="test", + title="title", + persona_id=None, + created_at=0, + updated_at=0, + ) + + class Manager: + async def get_conversation(self, **_kwargs): + return conversation + + service = ConversationService( + SimpleNamespace(), SimpleNamespace(conversation_manager=Manager()) + ) + exported = await service.export_conversations( + {"conversations": [{"user_id": "owner", "cid": "cid"}]} + ) + record = json.loads(exported.file_obj.read()) + image = record["content"][0]["content"][0]["image_url"] + assert image["detail"] == "high" + assert image["id"] == "img-1" diff --git a/tests/unit/test_image_lifecycle_workloads.py b/tests/unit/test_image_lifecycle_workloads.py new file mode 100644 index 0000000000..42427d35ef --- /dev/null +++ b/tests/unit/test_image_lifecycle_workloads.py @@ -0,0 +1,36 @@ +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path + + +def test_lifecycle_workload_defines_bounded_matrix_and_no_payload_dump(): + source = Path("scripts/image_memory_bench/lifecycle_workloads.py").read_text( + encoding="utf-8" + ) + assert "REQUESTS = 200" in source + assert "SESSIONS = 4" in source + assert "--request-count" in source + assert "--window-turns" in source + assert "time.sleep(0.01)" in source + assert "120" in source + assert "b64encode" not in source + assert "print(" not in source + + +def test_lifecycle_child_requires_real_database_and_reports_wire_metadata(tmp_path): + output = tmp_path / "result.jsonl" + result = subprocess.run( + [ + sys.executable, + "scripts/image_memory_bench/lifecycle_workloads.py", + str(output), + "--child", + ], + check=False, + capture_output=True, + text=True, + ) + assert result.returncode != 0 + assert not output.exists() diff --git a/tests/unit/test_image_media_store.py b/tests/unit/test_image_media_store.py new file mode 100644 index 0000000000..4adc030fe9 --- /dev/null +++ b/tests/unit/test_image_media_store.py @@ -0,0 +1,304 @@ +"""Tests for durable image object ownership and authorization.""" + +import io +import json +import os +from concurrent.futures import ThreadPoolExecutor + +import pytest +from PIL import Image + +from astrbot.core.utils.image_media_store import ImageMediaStore + + +def _png() -> bytes: + output = io.BytesIO() + with Image.new("RGBA", (9, 7), (20, 40, 60, 100)) as image: + image.save(output, "PNG") + return output.getvalue() + + +def test_store_deduplicates_and_round_trips_exact_bytes(tmp_path): + store = ImageMediaStore(tmp_path / "media") + data = _png() + first = store.put(data, detail="high") + second = store.put(data, detail="low") + + assert first.media_id == second.media_id + assert first.width == 9 and first.height == 7 + assert store.read(first, {first.media_id}) == data + assert len(list((tmp_path / "media").glob("*.bin"))) == 1 + + +def test_shared_blob_keeps_each_reference_metadata(tmp_path): + store = ImageMediaStore(tmp_path / "media") + data = _png() + first = store.put(data, detail="low", image_id="first") + second = store.put(data, detail="high", image_id="second") + assert first.detail == "low" and first.image_id == "first" + assert second.detail == "high" and second.image_id == "second" + + +def test_shared_blob_keeps_declared_mime_per_reference(tmp_path): + store = ImageMediaStore(tmp_path / "media") + data = _png() + first = store.put(data, mime_type="image/png") + second = store.put(data, mime_type="image/custom", image_id="custom") + assert first.mime_type == "image/png" + assert second.mime_type == "image/custom" + assert store.read(second, {second.media_id}) == data + + +def test_store_rejects_unauthorized_reference_and_missing_object(tmp_path): + store = ImageMediaStore(tmp_path / "media") + ref = store.put(_png()) + with pytest.raises(PermissionError): + store.read(ref, set()) + (tmp_path / "media" / f"{ref.media_id}.bin").unlink() + with pytest.raises(FileNotFoundError): + store.read(ref, {ref.media_id}) + + +def test_store_never_commits_invalid_or_partial_media(tmp_path): + store = ImageMediaStore(tmp_path / "media") + with pytest.raises(ValueError): + store.put(b"not an image") + assert not list((tmp_path / "media").glob("*.bin")) + assert not list((tmp_path / "media").glob("*.json")) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("runtime_message", [False, True]) +async def test_reference_materialization_preserves_history_and_bytes( + tmp_path, runtime_message +): + import base64 + + from astrbot.core.agent.message import Message + from astrbot.core.utils.image_media_store import materialize_image_media_refs + + store = ImageMediaStore(tmp_path / "media") + data = _png() + ref = store.put(data, detail="high") + history = [{"role": "user", "content": [ref.model_dump()]}] + messages = ( + [Message.model_validate(item) for item in history] + if runtime_message + else history + ) + result = await materialize_image_media_refs(messages, store) + dumped = result[0].model_dump() if runtime_message else result[0] + image = dumped["content"][0]["image_url"] + assert base64.b64decode(image["url"].split(",", 1)[1]) == data + assert image["detail"] == "high" + assert history[0]["content"][0]["type"] == "image_media_ref" + if runtime_message: + assert messages[0].content[0].type == "image_media_ref" + + +@pytest.mark.asyncio +async def test_unselected_reference_is_not_read(tmp_path, monkeypatch): + from astrbot.core.utils.image_media_store import materialize_image_media_refs + + store = ImageMediaStore(tmp_path / "media") + store.put(_png()) + + def fail_read(*args, **kwargs): + raise AssertionError("An unselected image must not be opened") + + monkeypatch.setattr(store, "read", fail_read) + selected = [{"role": "user", "content": "Only this text is in the active window"}] + assert await materialize_image_media_refs(selected, store) == selected + + +@pytest.mark.asyncio +async def test_missing_reference_has_bounded_placeholder(tmp_path): + from astrbot.core.utils.image_media_store import materialize_image_media_refs + + store = ImageMediaStore(tmp_path / "media") + ref = store.put(_png()) + (store.root / f"{ref.media_id}.bin").unlink() + result = await materialize_image_media_refs( + [{"role": "user", "content": [ref.model_dump()]}], store + ) + assert result[0]["content"] == [{"type": "text", "text": "[Image unavailable]"}] + + +def test_reference_rejects_path_traversal(): + from astrbot.core.utils.image_media_store import ImageMediaRef + + with pytest.raises(ValueError): + ImageMediaRef("../outside", "image/png", 1, 1, 10) + + +def test_deduplicated_corrupt_object_is_rejected(tmp_path): + store = ImageMediaStore(tmp_path / "media") + ref = store.put(_png()) + (store.root / f"{ref.media_id}.bin").write_bytes(b"corrupt") + with pytest.raises(OSError): + store.put(_png()) + + +@pytest.mark.parametrize("failure_point", ["link", "fsync"]) +def test_write_failure_leaves_no_usable_reference(tmp_path, monkeypatch, failure_point): + import astrbot.core.utils.image_media_store as media_store + + store = ImageMediaStore(tmp_path / "media") + original_link = os.link + original_fsync = os.fsync + + if failure_point == "link": + + def fail_link(source, target): + raise OSError("injected link failure") + + monkeypatch.setattr(media_store.os, "link", fail_link) + else: + + def fail_fsync(fd): + raise OSError("injected fsync failure") + + monkeypatch.setattr(media_store.os, "fsync", fail_fsync) + + with pytest.raises(OSError): + store.put(_png()) + assert not list(store.root.glob("*.json")) + if failure_point == "fsync": + assert not list(store.root.iterdir()) + assert original_link and original_fsync + + +def test_metadata_replace_failure_is_repaired_by_next_put(tmp_path, monkeypatch): + import astrbot.core.utils.image_media_store as media_store + + store = ImageMediaStore(tmp_path / "media") + original_link = os.link + failed = False + + def fail_metadata_link(source, target): + nonlocal failed + if target.suffix == ".json" and not failed: + failed = True + raise OSError("injected metadata link failure") + return original_link(source, target) + + monkeypatch.setattr(media_store.os, "link", fail_metadata_link) + with pytest.raises(OSError): + store.put(_png()) + assert not list(store.root.glob("*.json")) + ref = store.put(_png(), detail="next") + assert store.read(ref, {ref.media_id}) == _png() + + +def test_concurrent_same_bytes_puts_return_independent_references(tmp_path): + data = _png() + + def put(detail): + return ImageMediaStore(tmp_path / "media").put( + data, detail=detail, image_id=detail + ) + + with ThreadPoolExecutor(max_workers=2) as executor: + refs = list(executor.map(put, ["first", "second"])) + assert {ref.detail for ref in refs} == {"first", "second"} + assert {ref.image_id for ref in refs} == {"first", "second"} + for ref in refs: + assert ImageMediaStore(tmp_path / "media").read(ref, {ref.media_id}) == data + + +def test_partial_reference_metadata_missing_is_not_usable(tmp_path): + store = ImageMediaStore(tmp_path / "media") + data = _png() + ref = store.put(data) + (store.root / f"{ref.media_id}.json").unlink() + with pytest.raises(OSError, match="incomplete"): + store.read(ref, {ref.media_id}) + + +@pytest.mark.parametrize("suffix", [".bin", ".json"]) +def test_symlinked_media_entries_are_rejected(tmp_path, suffix): + store = ImageMediaStore(tmp_path / "media") + ref = store.put(_png()) + target = store.root / f"{ref.media_id}{suffix}" + target.unlink() + target.symlink_to(tmp_path / "outside") + with pytest.raises(OSError): + store.read(ref, {ref.media_id}) + + +def test_metadata_tampering_and_byte_size_tampering_are_rejected(tmp_path): + store = ImageMediaStore(tmp_path / "media") + ref = store.put(_png(), detail="high", image_id="original-id") + metadata_path = store.root / f"{ref.media_id}.json" + metadata = json.loads(metadata_path.read_text()) + metadata["detail"] = "low" + metadata_path.write_text(json.dumps(metadata)) + assert store.read(ref, {ref.media_id}) == _png() + metadata["detail"] = ref.detail + metadata["byte_size"] = ref.byte_size + 1 + metadata_path.write_text(json.dumps(metadata)) + with pytest.raises(OSError): + store.read(ref, {ref.media_id}) + + +@pytest.mark.asyncio +async def test_materialized_provider_prefix_preserves_mime_detail_id_and_bytes( + tmp_path, +): + import base64 + + from astrbot.core.utils.image_media_store import materialize_image_media_refs + + store = ImageMediaStore(tmp_path / "media") + data = _png() + ref = store.put(data, mime_type="image/png", detail="high", image_id="img-7") + result = await materialize_image_media_refs( + [{"role": "user", "content": [ref.model_dump()]}], store + ) + image_url = result[0]["content"][0]["image_url"] + assert image_url == { + "url": "data:image/png;base64," + base64.b64encode(data).decode(), + "detail": "high", + "id": "img-7", + } + + +def test_shared_reference_survives_restart_and_temp_cleanup(tmp_path): + data = _png() + first = ImageMediaStore(tmp_path / "media") + ref = first.put(data) + (first.root / "orphan.tmp").write_bytes(b"orphan") + restarted = ImageMediaStore(tmp_path / "media") + assert restarted.read(ref, {ref.media_id}) == data + assert (restarted.root / "orphan.tmp").exists() + + +def test_symlinked_store_root_is_rejected(tmp_path): + target = tmp_path / "target" + target.mkdir() + root = tmp_path / "media" + root.symlink_to(target, target_is_directory=True) + with pytest.raises(OSError, match="root"): + ImageMediaStore(root).put(_png()) + + +def test_memory_error_is_not_swallowed(tmp_path, monkeypatch): + store = ImageMediaStore(tmp_path / "media") + ref = store.put(_png()) + + def fail_read(*args, **kwargs): + raise MemoryError("injected") + + monkeypatch.setattr(store, "read", fail_read) + + async def run(): + from astrbot.core.utils.image_media_store import materialize_image_media_refs + + return await materialize_image_media_refs( + [{"role": "user", "content": [ref.model_dump()]}], store + ) + + with pytest.raises(MemoryError, match="injected"): + import asyncio + + asyncio.run(run()) diff --git a/tests/unit/test_image_memory_benchmark.py b/tests/unit/test_image_memory_benchmark.py new file mode 100644 index 0000000000..65ec479c97 --- /dev/null +++ b/tests/unit/test_image_memory_benchmark.py @@ -0,0 +1,59 @@ +"""Safety checks for the standalone image benchmark utilities.""" + +import json +import subprocess +import sys + + +def test_baseline_manifest_records_fixture_fingerprints(tmp_path): + fixture = tmp_path / "fixtures" / "7" + fixture.mkdir(parents=True) + (fixture / "sample.png").write_bytes(b"fixture") + output = tmp_path / "manifest.json" + subprocess.run( + [ + sys.executable, + "scripts/image_memory_bench/baseline_manifest.py", + str(output), + "--fixtures", + str(tmp_path / "fixtures"), + ], + check=True, + ) + manifest = json.loads(output.read_text()) + assert manifest["fixtures"][0]["sha256"] + + +def test_migration_dry_run_and_rollback_keep_input(tmp_path): + source = tmp_path / "history.jsonl" + source.write_text( + json.dumps({"history": [{"role": "user", "content": "hello"}]}) + "\n" + ) + output = tmp_path / "migrated.jsonl" + media = tmp_path / "media" + command = [ + sys.executable, + "scripts/image_memory_bench/migrate_history.py", + str(source), + str(output), + "--media-dir", + str(media), + ] + subprocess.run(command, check=True) + assert not output.exists() + subprocess.run(command + ["--apply"], check=True) + rollback = tmp_path / "rollback.jsonl" + subprocess.run( + [ + sys.executable, + "scripts/image_memory_bench/migrate_history.py", + str(output), + str(rollback), + "--media-dir", + str(media), + "--rollback", + "--apply", + ], + check=True, + ) + assert output.read_bytes() == rollback.read_bytes() diff --git a/tests/unit/test_image_preparation_budget.py b/tests/unit/test_image_preparation_budget.py new file mode 100644 index 0000000000..2abc322114 --- /dev/null +++ b/tests/unit/test_image_preparation_budget.py @@ -0,0 +1,376 @@ +"""Regression tests for bounded image preparation and output ownership.""" + +import asyncio +import base64 +import io +import random +import threading +from pathlib import Path + +import pytest +from PIL import Image + +from astrbot.core.utils import media_utils + + +def test_file_base64_encoding_stream_matches_standard_library(tmp_path): + source = tmp_path / "payload.bin" + payload = random.Random(31).randbytes(1023 * 1024 + 5) + source.write_bytes(payload) + + assert media_utils._encode_file_to_base64(source) == base64.b64encode( + payload + ).decode("ascii") + + +@pytest.mark.parametrize( + "size,target", [((1280, 1280), (960, 960)), ((131, 197), (71, 103))] +) +def test_strip_resize_matches_composited_lanczos(tmp_path, size, target): + from PIL import ImageChops + + with Image.frombytes( + "RGBA", size, random.Random(43).randbytes(size[0] * size[1] * 4) + ) as source: + with ( + source.resize(target, Image.Resampling.LANCZOS) as expected, + media_utils._resize_alpha_in_strips(source, target) as actual, + ): + assert actual.size == expected.size + for color in ("black", "white"): + with ( + Image.new("RGBA", target, color) as background, + Image.alpha_composite(background, expected) as expected_view, + Image.alpha_composite(background, actual) as actual_view, + ImageChops.difference(expected_view, actual_view) as delta, + ): + assert max(high for _, high in delta.getextrema()) <= 2 + + +@pytest.mark.parametrize( + "image_format,mode", [("PNG", "RGBA"), ("JPEG", "RGB"), ("WEBP", "RGBA")] +) +def test_compliant_bytes_are_preserved(tmp_path, image_format, mode): + path = tmp_path / "original" + with Image.new(mode, (32, 20)) as image: + image.save(path, image_format) + original = path.read_bytes() + result = media_utils._compress_image_sync(path, tmp_path, 1280, 95, True) + assert result is None + assert path.read_bytes() == original + assert list(tmp_path.iterdir()) == [path] + + +@pytest.mark.parametrize( + "mode,save_options,expected_format", + [("RGB", {}, "JPEG"), ("RGBA", {"lossless": True}, "PNG")], +) +def test_oversized_static_webp_avoids_pillow_decoder( + tmp_path, monkeypatch, mode, save_options, expected_format +): + if media_utils._get_webp_decoder() is None: + pytest.skip("Pillow does not expose the bundled WebP decoder") + + source = tmp_path / "source.webp" + pixel_count = 512 * 512 + channels = 4 if mode == "RGBA" else 3 + with Image.frombytes( + mode, + (512, 512), + random.Random(17).randbytes(pixel_count * channels), + ) as image: + image.save(source, "WEBP", **save_options) + + def unexpected_pillow_open(*args, **kwargs): + pytest.fail("Static WebP should use the preallocated decoder path") + + monkeypatch.setattr(media_utils.PILImage, "open", unexpected_pillow_open) + output = media_utils._compress_image_sync(source, tmp_path, 64, 95, True, 20_000) + monkeypatch.undo() + + assert output is not None + with Image.open(output) as prepared: + assert prepared.format == expected_format + assert max(prepared.size) <= 64 + + +def test_pixel_compliant_noise_fits_byte_budget(tmp_path): + source = tmp_path / "noise.png" + with Image.frombytes( + "RGBA", (1280, 1280), random.Random(7).randbytes(1280 * 1280 * 4) + ) as image: + image.save(source) + budget = 4 * 1024 * 1024 + assert 4 * ((source.stat().st_size + 2) // 3) > budget + result = media_utils._compress_image_sync(source, tmp_path, 1280, 95, True, budget) + assert result is not None + assert 4 * ((Path(result).stat().st_size + 2) // 3) <= budget + with Image.open(result) as image: + assert image.mode == "RGBA" + assert image.width < 1280 + + +def test_animation_is_preserved_when_compliant_and_rejected_when_oversized(tmp_path): + source = tmp_path / "animated.gif" + frames = [Image.new("RGB", (16, 12), color) for color in ("red", "blue")] + try: + frames[0].save( + source, + format="GIF", + save_all=True, + append_images=frames[1:], + duration=40, + loop=0, + ) + finally: + for frame in frames: + frame.close() + + original = source.read_bytes() + assert media_utils._compress_image_sync(source, tmp_path, 4, 95, True) is None + assert source.read_bytes() == original + + with pytest.raises(media_utils.ImagePayloadTooLargeError): + media_utils._compress_image_sync(source, tmp_path, 4, 95, True, 1) + assert list(tmp_path.glob("compressed_*")) == [] + + +def test_webp_animation_is_preserved_without_flattening(tmp_path): + source = tmp_path / "animated.webp" + frames = [Image.new("RGBA", (16, 12), color) for color in ("red", "blue")] + try: + frames[0].save( + source, + format="WEBP", + save_all=True, + append_images=frames[1:], + duration=40, + loop=0, + ) + finally: + for frame in frames: + frame.close() + + original = source.read_bytes() + assert media_utils._compress_image_sync(source, tmp_path, 4, 95, True) is None + assert source.read_bytes() == original + + with pytest.raises(media_utils.ImagePayloadTooLargeError): + media_utils._compress_image_sync(source, tmp_path, 4, 95, True, 1) + + +def test_screenshot_preserves_oriented_coordinates(tmp_path): + path = tmp_path / "oriented.jpg" + with Image.frombytes( + "RGB", (120, 80), random.Random(1).randbytes(120 * 80 * 3) + ) as image: + exif = image.getexif() + exif[274] = 6 + image.save(path, quality=100, exif=exif) + output = media_utils._compress_image_sync( + path, tmp_path, 10, 80, True, 10_000, preserve_dimensions=True + ) + assert output is not None + with Image.open(output) as image: + assert image.size == (80, 120) + assert image.getexif().get(274, 1) == 1 + + +@pytest.mark.parametrize("orientation", range(1, 9)) +def test_thumbnail_preserves_exif_corner_placement(tmp_path, orientation): + from PIL import ImageOps + + source = tmp_path / "corners.jpg" + with Image.new("RGB", (800, 600), "black") as image: + image.paste("red", (0, 0, 400, 300)) + image.paste("green", (400, 0, 800, 300)) + image.paste("blue", (0, 300, 400, 600)) + image.paste("white", (400, 300, 800, 600)) + exif = image.getexif() + exif[274] = orientation + image.save(source, quality=95, exif=exif) + with Image.open(source) as image: + expected = ImageOps.exif_transpose(image) + expected.thumbnail((160, 160)) + try: + output = media_utils._compress_image_sync(source, tmp_path, 160, 95, True) + assert output is not None + with Image.open(output) as actual: + assert actual.size == expected.size + assert actual.getexif().get(274, 1) == 1 + for x in (actual.width // 4, actual.width * 3 // 4): + for y in (actual.height // 4, actual.height * 3 // 4): + assert ( + max( + abs(a - b) + for a, b in zip( + actual.getpixel((x, y)), expected.getpixel((x, y)) + ) + ) + <= 5 + ) + finally: + expected.close() + + +def test_jpeg_stops_at_first_fitting_quality(tmp_path, monkeypatch): + path = tmp_path / "photo.jpg" + with Image.new("RGB", (200, 100), (20, 40, 80)) as image: + image.save(path, "JPEG", quality=95) + qualities = [] + original_save = Image.Image.save + + def save(image, target, fmt=None, **kwargs): + if fmt == "JPEG": + qualities.append(kwargs.get("quality")) + return original_save(image, target, fmt, **kwargs) + + monkeypatch.setattr(Image.Image, "save", save) + output = media_utils._compress_image_sync( + path, tmp_path, 100, 95, True, 4 * 1024 * 1024 + ) + assert output is not None + assert qualities == [95] + + +def test_cua_tries_next_jpeg_quality_without_resizing(tmp_path, monkeypatch): + path = tmp_path / "opaque.png" + with Image.frombytes( + "RGB", (120, 80), random.Random(23).randbytes(120 * 80 * 3) + ) as image: + image.save(path, "PNG") + qualities = [] + original_save = Image.Image.save + + def save(image, target, fmt=None, **kwargs): + if fmt == "JPEG": + quality = kwargs.get("quality") + qualities.append(quality) + if quality == 95: + original_save(image, target, fmt, **kwargs) + with Path(target).open("ab") as candidate: + candidate.write(b"x" * 20_000) + return + return original_save(image, target, fmt, **kwargs) + + monkeypatch.setattr(Image.Image, "save", save) + output = media_utils._compress_image_sync( + path, tmp_path, 1, 95, True, 20_000, preserve_dimensions=True + ) + assert output is not None + assert qualities[:2] == [95, 85] + with Image.open(output) as image: + assert image.size == (120, 80) + + +def test_impossible_screenshot_budget_rejects_without_resizing(tmp_path): + path = tmp_path / "screenshot.png" + with Image.new("RGBA", (100, 60), (20, 100, 50, 127)) as image: + image.save(path) + with pytest.raises(media_utils.ImagePayloadTooLargeError): + media_utils._compress_image_sync( + path, tmp_path, 1, 95, True, 1, preserve_dimensions=True + ) + assert list(tmp_path.iterdir()) == [path] + + +def test_write_failure_removes_partial_candidate(tmp_path, monkeypatch): + source = io.BytesIO() + with Image.new("RGB", (100, 60)) as image: + image.save(source, "PNG") + + def fail_write(image, path, *args, **kwargs): + Path(path).write_bytes(b"partial") + raise OSError("disk full") + + monkeypatch.setattr(Image.Image, "save", fail_write) + with pytest.raises(OSError, match="disk full"): + media_utils._compress_image_sync(source.getvalue(), tmp_path, 10, 95, True) + assert not list(tmp_path.iterdir()) + + +def test_decode_failure_leaves_no_output(tmp_path): + with pytest.raises(OSError): + media_utils._compress_image_sync(b"invalid", tmp_path, 10, 95, True) + assert not list(tmp_path.iterdir()) + + +@pytest.mark.asyncio +async def test_disabled_compression_rejects_oversize_before_read(tmp_path, monkeypatch): + source = tmp_path / "oversized.png" + with Image.new("RGB", (32, 32), "red") as image: + image.save(source) + + def unexpected_read(path): + pytest.fail("Oversized source was read before checking the byte limit") + + monkeypatch.setattr(Path, "read_bytes", unexpected_read) + with pytest.raises(media_utils.ImagePayloadTooLargeError): + await media_utils.prepare_image_source( + str(source), + options=media_utils.ImagePreparationOptions( + enabled=False, max_encoded_bytes=1 + ), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "failure", [media_utils.ImagePayloadTooLargeError, MemoryError, OSError, ValueError] +) +async def test_agent_does_not_fall_back_to_original_on_resource_error( + tmp_path, monkeypatch, failure +): + source = tmp_path / "image.png" + with Image.new("RGB", (8, 8), "red") as image: + image.save(source) + + async def fail(*args, **kwargs): + raise failure("bounded failure") + + monkeypatch.setattr(media_utils, "compress_image", fail) + with pytest.raises(failure): + await media_utils.prepare_image_source(str(source)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("timeout", [False, True]) +async def test_cancelled_worker_cleans_after_exit(tmp_path, monkeypatch, timeout): + source = tmp_path / "source" + source.write_bytes(b"source") + output = tmp_path / "worker-output" + entered = threading.Event() + release = threading.Event() + cleaned = asyncio.Event() + original_unlink = Path.unlink + + def unlink(path, *args, **kwargs): + original_unlink(path, *args, **kwargs) + if path == output: + cleaned.set() + + monkeypatch.setattr(Path, "unlink", unlink) + + def blocked_worker(*args, **kwargs): + entered.set() + release.wait(5) + output.write_bytes(b"prepared") + return str(output) + + monkeypatch.setattr(media_utils, "_compress_image_sync", blocked_worker) + task = asyncio.create_task(media_utils.compress_image(str(source))) + try: + assert await asyncio.to_thread(entered.wait, 2) + if timeout: + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait_for(task, timeout=0.01) + else: + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + finally: + release.set() + # A finished task can still have its cleanup callback queued on the loop. + await asyncio.wait_for(cleaned.wait(), timeout=5) + assert not output.exists() + assert source.read_bytes() == b"source" diff --git a/tests/unit/test_image_screenshot_encoding.py b/tests/unit/test_image_screenshot_encoding.py new file mode 100644 index 0000000000..cec79c5b0d --- /dev/null +++ b/tests/unit/test_image_screenshot_encoding.py @@ -0,0 +1,43 @@ +from pathlib import Path + +from PIL import Image, ImageDraw + +from astrbot.core.utils import media_utils + + +def test_prefers_smaller_lossless_encoding_for_ui_screenshot(tmp_path): + source = tmp_path / "screenshot-ordinary.png" + with Image.new("RGB", (1920, 1080), "white") as image: + draw = ImageDraw.Draw(image) + for x in range(0, 1920, 32): + draw.line((x, 0, x, 1080), fill=(205, 205, 205), width=1) + for y in range(0, 1080, 32): + draw.line((0, y, 1920, y), fill=(205, 205, 205), width=1) + for index in range(80): + draw.text( + (20 + (index % 10) * 185, 20 + (index // 10) * 125), + f"UI {index:02d}", + fill="black", + ) + image.save(source, "PNG", optimize=True) + + output = media_utils._compress_image_sync( + source, tmp_path, 1280, 95, True, 4 * 1024 * 1024 + ) + + assert output is not None + with Image.open(source) as original: + resized = original.copy() + resized.thumbnail((1280, 1280), Image.Resampling.LANCZOS) + png_candidate = tmp_path / "expected.png" + jpeg_candidate = tmp_path / "expected.jpg" + resized.save(png_candidate, "PNG", optimize=True) + resized.save(jpeg_candidate, "JPEG", quality=95, optimize=True) + # The old implementation always selected JPEG for opaque PNG input, + # even when the same resized pixels have a smaller PNG encoding. + assert png_candidate.stat().st_size < jpeg_candidate.stat().st_size + assert max(resized.size) <= 1280 + assert Path(output).suffix == ".png" + assert Path(output).stat().st_size == png_candidate.stat().st_size + with Image.open(output) as image: + assert max(image.size) <= 1280 diff --git a/tests/unit/test_image_source_integration.py b/tests/unit/test_image_source_integration.py new file mode 100644 index 0000000000..41623a78d8 --- /dev/null +++ b/tests/unit/test_image_source_integration.py @@ -0,0 +1,116 @@ +import base64 +import io + +import pytest +from PIL import Image + +from astrbot.core.agent.message import ImageURLPart +from astrbot.core.provider.entities import ProviderRequest +from astrbot.core.utils.image_media_store import ( + ImageMediaStore, + persist_inline_image_refs, +) +from astrbot.core.utils.media_utils import ( + ImagePayloadTooLargeError, + ImagePreparationInput, + ImagePreparationOptions, + prepare_image_source, +) + + +def _image(size=(40, 20), image_format="PNG"): + output = io.BytesIO() + with Image.new("RGB", size, (30, 60, 90)) as image: + image.save(output, image_format) + return output.getvalue() + + +@pytest.mark.asyncio +async def test_current_plugin_image_part_is_prepared_without_mutating_caller(tmp_path): + source = tmp_path / "quoted.png" + original = _image() + source.write_bytes(original) + part = ImageURLPart( + image_url=ImageURLPart.ImageURL( + url=str(source), id="plugin-image", detail="high" + ) + ).mark_as_temp() + + context = await ProviderRequest(extra_user_content_parts=[part]).assemble_context() + payload = context["content"][0] + + assert payload["image_url"]["id"] == "plugin-image" + assert payload["image_url"]["detail"] == "high" + assert payload["_no_save"] is True + assert payload["image_url"]["url"].startswith("data:image/png;base64,") + assert part.image_url.url == str(source) + + +@pytest.mark.asyncio +async def test_quoted_user_path_uses_configured_encoded_limit(tmp_path): + source = tmp_path / "quoted.png" + source.write_bytes(_image()) + options = ImagePreparationOptions(max_encoded_bytes=1) + + with pytest.raises(ImagePayloadTooLargeError): + await prepare_image_source(str(source), options=options) + + +@pytest.mark.asyncio +async def test_cua_dimensions_are_explicit_and_normal_tool_can_resize(tmp_path): + source = tmp_path / "tool.png" + source.write_bytes(_image((80, 40))) + normal = await prepare_image_source( + str(source), options=ImagePreparationOptions(max_size=20) + ) + cua = await prepare_image_source( + str(source), + options=ImagePreparationOptions(max_size=20, preserve_dimensions=True), + ) + + with Image.open(io.BytesIO(normal.to_bytes())) as image: + assert max(image.size) <= 20 + with Image.open(io.BytesIO(cua.to_bytes())) as image: + assert image.size == (80, 40) + + +def test_temporary_image_part_is_not_persisted(tmp_path): + image_bytes = _image() + image_url = "data:image/png;base64," + base64.b64encode(image_bytes).decode("ascii") + history = [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": image_url}, + "_no_save": True, + } + ], + } + ] + + stored = persist_inline_image_refs(history, ImageMediaStore(tmp_path / "media")) + + assert stored == history + assert not (tmp_path / "media").exists() + + +@pytest.mark.asyncio +async def test_shared_preparation_input_releases_owned_source(tmp_path): + source = tmp_path / "plugin-source.png" + owned = tmp_path / "resolver-owned.tmp" + data = _image((32, 24)) + source.write_bytes(data) + owned.write_bytes(b"temporary source") + + prepared = await prepare_image_source( + ImagePreparationInput( + str(source), + source_kind="plugin_mcp", + cleanup_paths=(owned,), + ) + ) + + assert prepared.to_bytes() == data + assert not owned.exists() diff --git a/tests/unit/test_image_source_preparation.py b/tests/unit/test_image_source_preparation.py new file mode 100644 index 0000000000..50c46a1c61 --- /dev/null +++ b/tests/unit/test_image_source_preparation.py @@ -0,0 +1,134 @@ +import asyncio +import base64 +import io +import threading +from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path + +import pytest +from PIL import Image + +from astrbot.core.agent.message import ImageURLPart +from astrbot.core.provider.entities import ProviderRequest +from astrbot.core.utils import media_utils + + +def _png() -> bytes: + stream = io.BytesIO() + with Image.new("RGBA", (8, 6), (20, 40, 60, 128)) as image: + image.save(stream, "PNG") + return stream.getvalue() + + +@pytest.mark.asyncio +async def test_prepare_image_source_accepts_all_reference_forms(tmp_path, monkeypatch): + data = _png() + source = tmp_path / "image.png" + source.write_bytes(data) + monkeypatch.chdir(tmp_path) + + class Handler(SimpleHTTPRequestHandler): + def log_message(self, *_args): + pass + + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + encoded = base64.b64encode(data).decode() + refs = [ + str(source), + f"http://127.0.0.1:{server.server_port}/image.png", + f"data:image/png;base64,{encoded}", + f"base64://{encoded}", + encoded, + ] + for ref in refs: + result = await media_utils.prepare_image_source(ref) + assert result.mime_type == "image/png" + assert result.to_bytes() == data + finally: + server.shutdown() + server.server_close() + + +@pytest.mark.asyncio +async def test_prepare_cancel_keeps_resolver_source_until_worker_exits( + tmp_path, monkeypatch +): + entered = threading.Event() + release = threading.Event() + source_seen = {} + original = media_utils._compress_image_sync + + def blocked(source, *args, **kwargs): + source_seen["path"] = Path(source) if isinstance(source, (str, Path)) else None + entered.set() + release.wait(5) + return original(source, *args, **kwargs) + + monkeypatch.setattr(media_utils, "_compress_image_sync", blocked) + encoded = base64.b64encode(_png()).decode() + task = asyncio.create_task( + media_utils.prepare_image_source(f"data:image/png;base64,{encoded}") + ) + assert await asyncio.to_thread(entered.wait, 2) + assert source_seen["path"] is not None + assert source_seen["path"].exists() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert source_seen["path"].exists() + release.set() + await asyncio.sleep(0.2) + assert not source_seen["path"].exists() + + +@pytest.mark.asyncio +async def test_prepare_bytes_write_failure_cleans_owned_source(monkeypatch, tmp_path): + owned = tmp_path / "owned.bin" + + monkeypatch.setattr(media_utils, "_temp_media_path", lambda *_args: owned) + + original_write_bytes = Path.write_bytes + + def fail_write(path, data): + original_write_bytes(path, b"partial") + raise OSError("disk full") + + monkeypatch.setattr(Path, "write_bytes", fail_write) + with pytest.raises(OSError, match="disk full"): + await media_utils.prepare_image_source(_png()) + assert not owned.exists() + + +@pytest.mark.asyncio +async def test_provider_request_passes_preparation_options(): + data = base64.b64encode(_png()).decode() + request = ProviderRequest( + image_urls=[f"data:image/png;base64,{data}"], + image_preparation_options=media_utils.ImagePreparationOptions( + enabled=False, max_encoded_bytes=None + ), + ) + context = await request.assemble_context() + assert context["content"][1]["image_url"]["url"].endswith(data) + + +@pytest.mark.asyncio +async def test_extra_image_part_prepares_copy_and_preserves_metadata(tmp_path): + data = _png() + source = tmp_path / "plugin.png" + source.write_bytes(data) + part = ImageURLPart( + image_url=ImageURLPart.ImageURL(url=str(source), id="plugin-id", detail="high") + ).mark_as_temp() + request = ProviderRequest(extra_user_content_parts=[part]) + + context = await request.assemble_context() + payload = context["content"][0] + assert payload["image_url"]["id"] == "plugin-id" + assert payload["image_url"]["detail"] == "high" + assert payload["_no_save"] is True + assert payload["image_url"]["url"].startswith("data:image/png;base64,") + assert part.image_url.url == str(source) diff --git a/tests/unit/test_manage_image_history.py b/tests/unit/test_manage_image_history.py new file mode 100644 index 0000000000..959db2f520 --- /dev/null +++ b/tests/unit/test_manage_image_history.py @@ -0,0 +1,354 @@ +from __future__ import annotations + +import base64 +import importlib.util +import json +import sqlite3 +import subprocess +import sys + +import pytest +from PIL import Image + + +def _image() -> bytes: + import io + + output = io.BytesIO() + Image.new("RGB", (3, 2), "red").save(output, "PNG") + return output.getvalue() + + +def _db(path, content): + connection = sqlite3.connect(path) + connection.execute( + "CREATE TABLE conversations (inner_conversation_id INTEGER PRIMARY KEY, conversation_id TEXT, content JSON)" + ) + connection.execute( + "INSERT INTO conversations VALUES (1, 'one', ?)", (json.dumps(content),) + ) + connection.commit() + connection.close() + + +def _run(db, media, *args): + return subprocess.run( + [ + sys.executable, + "scripts/manage_image_history.py", + str(db), + "--media", + str(media), + *args, + ], + check=False, + capture_output=True, + text=True, + ) + + +def test_migration_defaults_to_dry_run(tmp_path): + data = _image() + history = [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64," + + base64.b64encode(data).decode() + }, + } + ], + } + ] + db = tmp_path / "db.sqlite" + media = tmp_path / "media" + _db(db, history) + result = _run(db, media) + assert result.returncode == 0 + assert not media.exists() + assert ( + json.loads( + sqlite3.connect(db) + .execute("SELECT content FROM conversations") + .fetchone()[0] + ) + == history + ) + + +def test_apply_creates_backup_and_references(tmp_path): + data = _image() + history = [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64," + + base64.b64encode(data).decode(), + "id": "keep", + }, + } + ], + } + ] + db = tmp_path / "db.sqlite" + media = tmp_path / "media" + backup = tmp_path / "backup" + _db(db, history) + result = _run(db, media, "--apply", "--offline", "--backup", str(backup)) + assert result.returncode == 0, result.stderr + content = json.loads( + sqlite3.connect(db).execute("SELECT content FROM conversations").fetchone()[0] + ) + ref = content[0]["content"][0] + assert ref["type"] == "image_media_ref" and ref["image_id"] == "keep" + assert (backup / "db.sqlite").exists() and (backup / "manifest.json").exists() + + +def test_apply_requires_offline_and_cleanup_dry_run_is_read_only(tmp_path): + db = tmp_path / "db.sqlite" + media = tmp_path / "media" + _db(db, []) + media.mkdir() + (media / "dead.bin").write_bytes(b"x") + result = _run(db, media, "--apply") + assert result.returncode != 0 + result = _run(db, media, "--cleanup") + assert result.returncode == 0 + assert (media / "dead.bin").exists() + + +def test_cleanup_aborts_on_malformed_history(tmp_path): + db = tmp_path / "db.sqlite" + media = tmp_path / "media" + _db(db, {"not": "a history"}) + media.mkdir() + (media / ("a" * 64 + ".bin")).write_bytes(b"orphan") + result = _run(db, media, "--cleanup", "--apply", "--offline") + assert result.returncode != 0 + assert (media / ("a" * 64 + ".bin")).exists() + + +def test_cleanup_preserves_references_nested_in_retained_messages(tmp_path): + from astrbot.core.utils.image_media_store import ImageMediaStore + + media = tmp_path / "media" + store = ImageMediaStore(media) + ref = store.put(_image()) + db = tmp_path / "db.sqlite" + _db(db, [{"role": "tool", "content": [{"resource": ref.model_dump()}]}]) + result = _run(db, media, "--cleanup", "--apply", "--offline") + assert result.returncode == 0, result.stderr + assert store.read(ref, {ref.media_id}) == _image() + + +def test_cleanup_keeps_shared_media_until_last_history_reference_is_removed(tmp_path): + from astrbot.core.utils.image_media_store import ImageMediaStore + + media = tmp_path / "media" + store = ImageMediaStore(media) + ref = store.put(_image()) + first_history = [{"role": "user", "content": [ref.model_dump()]}] + second_history = [{"role": "user", "content": [ref.model_dump()]}] + db = tmp_path / "db.sqlite" + _db(db, first_history) + connection = sqlite3.connect(db) + connection.execute( + "INSERT INTO conversations VALUES (2, 'two', ?)", + (json.dumps(second_history),), + ) + connection.commit() + connection.close() + + first_delete = _run(db, media, "--cleanup", "--apply", "--offline") + assert first_delete.returncode == 0, first_delete.stderr + assert (media / f"{ref.media_id}.bin").exists() + + connection = sqlite3.connect(db) + connection.execute("DELETE FROM conversations WHERE inner_conversation_id=1") + connection.commit() + connection.close() + second_delete = _run(db, media, "--cleanup", "--apply", "--offline") + assert second_delete.returncode == 0, second_delete.stderr + assert (media / f"{ref.media_id}.bin").exists() + + connection = sqlite3.connect(db) + connection.execute("DELETE FROM conversations WHERE inner_conversation_id=2") + connection.commit() + connection.close() + last_delete = _run(db, media, "--cleanup", "--apply", "--offline") + assert last_delete.returncode == 0, last_delete.stderr + assert not (media / f"{ref.media_id}.bin").exists() + quarantine = media.parent / "media.quarantine" + assert (quarantine / f"{ref.media_id}.bin").exists() + + +def test_cleanup_leaves_unpaired_and_symlink_objects(tmp_path): + db = tmp_path / "db.sqlite" + media = tmp_path / "media" + _db(db, []) + media.mkdir() + media_id = "b" * 64 + (media / f"{media_id}.bin").write_bytes(b"bad") + (media / f"{media_id}.json").write_text("{}") + (media / ("c" * 64 + ".bin")).write_bytes(b"bad") + result = _run(db, media, "--cleanup", "--apply", "--offline") + assert result.returncode == 0 + assert (media / f"{media_id}.bin").exists() + + +def test_restore_refuses_newer_history_and_wal_is_not_left_stale(tmp_path): + data = _image() + history = [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64," + + base64.b64encode(data).decode() + }, + } + ], + } + ] + db = tmp_path / "db.sqlite" + media = tmp_path / "media" + backup = tmp_path / "backup" + _db(db, history) + assert ( + _run(db, media, "--apply", "--offline", "--backup", str(backup)).returncode == 0 + ) + connection = sqlite3.connect(db) + connection.execute("UPDATE conversations SET content='[]'") + connection.execute("CREATE TABLE unrelated (value TEXT)") + connection.execute("INSERT INTO unrelated VALUES ('changed')") + connection.commit() + connection.close() + (db.with_name(db.name + "-wal")).write_bytes(b"stale") + result = _run(db, media, "--restore", str(backup), "--offline") + assert result.returncode != 0 + assert not (db.with_name(db.name + ".before-restore")).exists() + + +def test_restore_validates_backup_and_conflicting_media_before_database_write(tmp_path): + db = tmp_path / "db.sqlite" + media = tmp_path / "media" + backup = tmp_path / "backup" + _db(db, []) + assert ( + _run(db, media, "--apply", "--offline", "--backup", str(backup)).returncode == 0 + ) + manifest = backup / "manifest.json" + manifest.write_text( + manifest.read_text().replace( + '"database":', '"database": "tampered", "ignored":' + ) + ) + before = db.read_bytes() + result = _run(db, media, "--restore", str(backup), "--offline") + assert result.returncode != 0 + assert db.read_bytes() == before + + +def test_successful_restore_preserves_inline_history_and_detail(tmp_path): + data = _image() + history = [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64," + + base64.b64encode(data).decode(), + "detail": "high", + }, + } + ], + } + ] + db = tmp_path / "db.sqlite" + media = tmp_path / "media" + backup = tmp_path / "backup" + _db(db, history) + assert ( + _run(db, media, "--apply", "--offline", "--backup", str(backup)).returncode == 0 + ) + result = _run(db, media, "--restore", str(backup), "--offline") + assert result.returncode == 0, result.stderr + restored = json.loads( + sqlite3.connect(db).execute("SELECT content FROM conversations").fetchone()[0] + ) + assert restored == history + + +def test_backup_reads_consistent_open_wal_database(tmp_path): + db = tmp_path / "db.sqlite" + media = tmp_path / "media" + backup = tmp_path / "backup" + connection = sqlite3.connect(db) + connection.execute("PRAGMA journal_mode=WAL") + connection.execute( + "CREATE TABLE conversations (inner_conversation_id INTEGER PRIMARY KEY, conversation_id TEXT, content JSON)" + ) + connection.execute("INSERT INTO conversations VALUES (1, 'wal', '[]')") + connection.commit() + assert (db.with_name(db.name + "-wal")).exists() + connection.close() + result = _run(db, media, "--apply", "--offline", "--backup", str(backup)) + assert result.returncode == 0, result.stderr + assert ( + sqlite3.connect(backup / "db.sqlite") + .execute("SELECT conversation_id FROM conversations") + .fetchone()[0] + == "wal" + ) + + +def test_restore_replace_failure_leaves_database_and_cleans_stage( + tmp_path, monkeypatch +): + db = tmp_path / "db.sqlite" + media = tmp_path / "media" + backup = tmp_path / "backup" + _db(db, []) + from astrbot.core.utils.image_media_store import ImageMediaStore + + ImageMediaStore(media).put(_image(), detail="high") + assert ( + _run(db, media, "--apply", "--offline", "--backup", str(backup)).returncode == 0 + ) + for path in media.iterdir(): + path.unlink() + module_spec = importlib.util.spec_from_file_location( + "manage_image_history", "scripts/manage_image_history.py" + ) + module = importlib.util.module_from_spec(module_spec) + module_spec.loader.exec_module(module) + + def fail_replace(self, target): + raise OSError("injected replace failure") + + monkeypatch.setattr(module.Path, "replace", fail_replace) + args = type( + "Args", + (), + { + "offline": True, + "restore": str(backup), + "database": str(db), + "media": str(media), + }, + ) + before = db.read_bytes() + with pytest.raises(OSError, match="injected replace failure"): + module._restore(args) + assert db.read_bytes() == before + assert not list(tmp_path.glob("astrbot-restore-media-*")) diff --git a/tests/unit/test_provider_image_references.py b/tests/unit/test_provider_image_references.py new file mode 100644 index 0000000000..0089dc9f36 --- /dev/null +++ b/tests/unit/test_provider_image_references.py @@ -0,0 +1,191 @@ +import base64 +from copy import deepcopy + +import pytest +from PIL import Image + +from astrbot.core.provider.sources.anthropic_source import ProviderAnthropic +from astrbot.core.provider.sources.gemini_source import ProviderGoogleGenAI +from astrbot.core.provider.sources.openai_source import ProviderOpenAIOfficial +from astrbot.core.utils.image_media_store import ImageMediaStore +from astrbot.core.utils.media_utils import ImagePayloadTooLargeError + + +@pytest.fixture +def image_context(tmp_path, monkeypatch): + with Image.new("RGB", (4, 3), "red") as image: + source = tmp_path / "source.png" + image.save(source, "PNG") + ref = ImageMediaStore(tmp_path / "media").put(source.read_bytes()) + context = [{"role": "user", "content": [ref.model_dump()]}] + for module in ("openai_source", "anthropic_source", "gemini_source"): + monkeypatch.setattr( + f"astrbot.core.provider.sources.{module}.get_astrbot_data_path", + lambda tmp_path=tmp_path: str(tmp_path), + ) + return context + + +def _bare(adapter): + instance = object.__new__(adapter) + instance.provider_config = {} + instance.provider_settings = {} + instance.model_name = "test-model" + instance.client = type( + "Client", (), {"base_url": type("URL", (), {"host": ""})()} + )() + return instance + + +@pytest.mark.asyncio +async def test_openai_payload_materializes_reference_without_mutating(image_context): + adapter = _bare(ProviderOpenAIOfficial) + payload, _ = await adapter._prepare_chat_payload(None, contexts=image_context) + assert payload["messages"][0]["content"][0]["image_url"]["url"].startswith( + "data:image/" + ) + assert image_context[0]["content"][0]["type"] == "image_media_ref" + + +@pytest.mark.asyncio +async def test_oversized_reference_fails_before_materialization(image_context): + oversized = deepcopy(image_context) + oversized[0]["content"][0]["byte_size"] = 10_000_000 + adapter = _bare(ProviderOpenAIOfficial) + with pytest.raises(ImagePayloadTooLargeError): + await adapter._prepare_chat_payload(None, contexts=oversized) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("adapter_type", [ProviderOpenAIOfficial, ProviderGoogleGenAI]) +async def test_provider_settings_control_reference_budget(adapter_type, image_context): + adapter = _bare(adapter_type) + adapter.provider_settings = {"image_compress_options": {"max_encoded_bytes": 1}} + with pytest.raises(ImagePayloadTooLargeError): + if adapter_type is ProviderOpenAIOfficial: + await adapter._prepare_chat_payload(None, contexts=image_context) + else: + await adapter._prepare_conversation({"messages": image_context}) + + adapter.provider_settings = { + "image_compress_options": {"max_encoded_bytes": 1024 * 1024} + } + if adapter_type is ProviderOpenAIOfficial: + payload, _ = await adapter._prepare_chat_payload(None, contexts=image_context) + assert payload["messages"][0]["content"][0]["image_url"] + else: + contents = await adapter._prepare_conversation({"messages": image_context}) + assert contents[0].parts[0].inline_data is not None + + +@pytest.mark.asyncio +async def test_text_only_modality_scrubs_missing_oversized_reference_without_read( + image_context, +): + missing = deepcopy(image_context) + missing[0]["content"][0]["media_id"] = "f" * 64 + missing[0]["content"][0]["byte_size"] = 100_000_000 + adapter = _bare(ProviderOpenAIOfficial) + adapter.provider_config = {"modalities": ["text"]} + payload, _ = await adapter._prepare_chat_payload(None, contexts=missing) + assert payload["messages"][0]["content"] == [{"type": "text", "text": "[Image]"}] + + +@pytest.mark.asyncio +async def test_anthropic_text_chat_materializes_reference_without_mutating( + image_context, +): + adapter = _bare(ProviderAnthropic) + captured = {} + + async def query(payloads, tools, **kwargs): + captured.update(payloads) + raise RuntimeError("stop after capture") + + adapter._query = query + with pytest.raises(RuntimeError, match="stop after capture"): + await adapter.text_chat(contexts=deepcopy(image_context)) + image = captured["messages"][0]["content"][0] + assert image["source"]["type"] == "base64" + assert image_context[0]["content"][0]["type"] == "image_media_ref" + + +@pytest.mark.asyncio +async def test_anthropic_provider_settings_control_reference_budget(image_context): + adapter = _bare(ProviderAnthropic) + adapter.provider_settings = {"image_compress_options": {"max_encoded_bytes": 1}} + with pytest.raises(ImagePayloadTooLargeError): + await adapter.text_chat(contexts=image_context) + + +@pytest.mark.asyncio +async def test_gemini_conversation_materializes_reference(image_context): + adapter = _bare(ProviderGoogleGenAI) + contents = await adapter._prepare_conversation({"messages": image_context}) + assert contents[0].parts[0].inline_data is not None + assert image_context[0]["content"][0]["type"] == "image_media_ref" + + +def test_anthropic_payload_detects_data_uri_from_prefix_without_copying_payload( + monkeypatch, +): + adapter = _bare(ProviderAnthropic) + image_bytes = b"\x89PNG\r\n\x1a\n" + b"x" * 4096 + encoded = base64.b64encode(image_bytes).decode() + detected_prefixes = [] + + def detect_mime(prefix: bytes) -> str: + detected_prefixes.append(prefix) + return "image/png" + + monkeypatch.setattr(adapter, "_detect_image_mime_type", detect_mime) + _, messages = adapter._prepare_payload( + [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": f"data:image/png;base64,{encoded}"}, + } + ], + } + ] + ) + + assert detected_prefixes == [image_bytes[:48]] + assert messages[0]["content"][0]["source"]["data"] == encoded + + +@pytest.mark.asyncio +async def test_gemini_data_uri_skips_shared_resolver(monkeypatch): + adapter = _bare(ProviderGoogleGenAI) + image_bytes = b"\x89PNG\r\n\x1a\n" + b"payload" + encoded = base64.b64encode(image_bytes).decode() + + async def fail_if_resolved(*_args, **_kwargs): + raise AssertionError("data URLs should not go through the shared resolver") + + monkeypatch.setattr( + "astrbot.core.provider.sources.gemini_source.resolve_media_ref_to_base64_data", + fail_if_resolved, + ) + contents = await adapter._prepare_conversation( + { + "messages": [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": f"data:image/png;base64,{encoded}"}, + } + ], + } + ] + } + ) + + assert contents[0].parts is not None + assert contents[0].parts[0].inline_data.data == image_bytes + assert contents[0].parts[0].inline_data.mime_type == "image/png" diff --git a/tests/unit/test_provider_request_retry.py b/tests/unit/test_provider_request_retry.py new file mode 100644 index 0000000000..95e3fc8832 --- /dev/null +++ b/tests/unit/test_provider_request_retry.py @@ -0,0 +1,40 @@ +import pytest + +from astrbot.core.exceptions import ProviderRequestTooLargeError +from astrbot.core.provider.sources.request_retry import retry_provider_request + + +class _Response: + status_code = 413 + + +class _RequestTooLargeError(Exception): + response = _Response() + + +@pytest.mark.asyncio +async def test_memory_error_is_not_retried(): + calls = 0 + + async def request(): + nonlocal calls + calls += 1 + raise MemoryError("allocation failed") + + with pytest.raises(MemoryError): + await retry_provider_request("test", request, max_attempts=5) + assert calls == 1 + + +@pytest.mark.asyncio +async def test_http_413_becomes_typed_request_size_error_without_retry(): + calls = 0 + + async def request(): + nonlocal calls + calls += 1 + raise _RequestTooLargeError("payload too large") + + with pytest.raises(ProviderRequestTooLargeError, match="HTTP 413"): + await retry_provider_request("test", request, max_attempts=5) + assert calls == 1 From 17169a673fc843f07716850daef929d83097e04c Mon Sep 17 00:00:00 2001 From: zenfun Date: Thu, 17 Sep 2026 18:22:55 +0800 Subject: [PATCH 5/7] chore: remove image memory experiment artifacts --- docs/en/dev/openapi-scopes.md | 1 - docs/public/openapi.json | 49 ------ docs/zh/dev/openapi-scopes.md | 1 - .../.openspec.yaml | 2 - .../optimize-image-memory-lifecycle/design.md | 152 ------------------ .../proposal.md | 28 ---- .../specs/agent-context-image-budget/spec.md | 40 ----- .../specs/conversation-history-media/spec.md | 37 ----- .../specs/image-memory-lifecycle/spec.md | 26 --- .../optimize-image-memory-lifecycle/tasks.md | 36 ----- 10 files changed, 372 deletions(-) delete mode 100644 openspec/changes/optimize-image-memory-lifecycle/.openspec.yaml delete mode 100644 openspec/changes/optimize-image-memory-lifecycle/design.md delete mode 100644 openspec/changes/optimize-image-memory-lifecycle/proposal.md delete mode 100644 openspec/changes/optimize-image-memory-lifecycle/specs/agent-context-image-budget/spec.md delete mode 100644 openspec/changes/optimize-image-memory-lifecycle/specs/conversation-history-media/spec.md delete mode 100644 openspec/changes/optimize-image-memory-lifecycle/specs/image-memory-lifecycle/spec.md delete mode 100644 openspec/changes/optimize-image-memory-lifecycle/tasks.md diff --git a/docs/en/dev/openapi-scopes.md b/docs/en/dev/openapi-scopes.md index 7b44c6e9b7..4bf87c6c4b 100644 --- a/docs/en/dev/openapi-scopes.md +++ b/docs/en/dev/openapi-scopes.md @@ -184,7 +184,6 @@ Manage conversations and platform-session data. | `GET` | `/api/v1/conversations/{conversation_id}` | — | | `PATCH` | `/api/v1/conversations/{conversation_id}` | — | | `DELETE` | `/api/v1/conversations/{conversation_id}` | — | -| `GET` | `/api/v1/conversations/{conversation_id}/media/{media_id}` | — | | `PUT` | `/api/v1/conversations/{conversation_id}/messages` | — | | `GET` | `/api/v1/session-groups` | — | | `POST` | `/api/v1/session-groups` | — | diff --git a/docs/public/openapi.json b/docs/public/openapi.json index 6a9d4a5f2c..b666e92afa 100644 --- a/docs/public/openapi.json +++ b/docs/public/openapi.json @@ -6457,55 +6457,6 @@ "description": "**Required scope:** `data`" } }, - "/api/v1/conversations/{conversation_id}/media/{media_id}": { - "get": { - "tags": [ - "Conversations" - ], - "summary": "Preview an image referenced by a conversation", - "operationId": "previewConversationMedia", - "x-astrbot-scope": "data", - "parameters": [ - { - "$ref": "#/components/parameters/ConversationId" - }, - { - "name": "media_id", - "in": "path", - "required": true, - "schema": { - "type": "string", - "pattern": "^[0-9a-f]{64}$" - } - }, - { - "name": "user_id", - "in": "query", - "required": true, - "schema": { - "type": "string" - } - } - ], - "responses": { - "200": { - "description": "Image bytes", - "content": { - "image/*": { - "schema": { - "type": "string", - "format": "binary" - } - } - } - }, - "404": { - "description": "Conversation or media is unavailable" - } - }, - "description": "**Required scope:** `data`" - } - }, "/api/v1/conversations/export": { "post": { "tags": [ diff --git a/docs/zh/dev/openapi-scopes.md b/docs/zh/dev/openapi-scopes.md index 36146c7ef8..3152454eb3 100644 --- a/docs/zh/dev/openapi-scopes.md +++ b/docs/zh/dev/openapi-scopes.md @@ -184,7 +184,6 @@ outline: deep | `GET` | `/api/v1/conversations/{conversation_id}` | — | | `PATCH` | `/api/v1/conversations/{conversation_id}` | — | | `DELETE` | `/api/v1/conversations/{conversation_id}` | — | -| `GET` | `/api/v1/conversations/{conversation_id}/media/{media_id}` | — | | `PUT` | `/api/v1/conversations/{conversation_id}/messages` | — | | `GET` | `/api/v1/session-groups` | — | | `POST` | `/api/v1/session-groups` | — | diff --git a/openspec/changes/optimize-image-memory-lifecycle/.openspec.yaml b/openspec/changes/optimize-image-memory-lifecycle/.openspec.yaml deleted file mode 100644 index 96db9a43b6..0000000000 --- a/openspec/changes/optimize-image-memory-lifecycle/.openspec.yaml +++ /dev/null @@ -1,2 +0,0 @@ -schema: spec-driven -created: 2026-09-15 diff --git a/openspec/changes/optimize-image-memory-lifecycle/design.md b/openspec/changes/optimize-image-memory-lifecycle/design.md deleted file mode 100644 index 47d11ea4a0..0000000000 --- a/openspec/changes/optimize-image-memory-lifecycle/design.md +++ /dev/null @@ -1,152 +0,0 @@ -## Context - -See `proposal.md` and the three delta specs. The current provider request path turns image references into complete data URIs, while conversation saving serializes the model-visible message list. Image preparation currently focuses on longest-edge pixels and can pass through large pixel-compliant files. Existing histories already contain inline data URIs and must remain readable. - -## Goals / Non-Goals - -**Goals:** - -- Establish one preparation boundary shared by all image-producing inputs. -- Bound the encoded payload that can enter a provider request. -- Separate durable conversation records from media bytes and resolve media only for selected active messages. -- Keep old inline histories readable and make migration explicit. -- Measure file size, encoded payload size, request size, RSS high-water mark, Python allocations, temporary files, and post-request retention. - -**Non-Goals:** - -- Do not silently remove images from the active context merely to improve memory metrics. -- Do not rewrite existing conversations during ordinary reads. -- Do not assume the 5 MiB limit from #10089 applies to every provider. -- Do not make the provider-specific wire format depend on the persistence representation. - -## Decisions - -### 1. Use a media reference as the persistence boundary - -New history image parts use a versioned internal media reference containing a content hash, MIME type, dimensions, encoded byte size, and original image-detail metadata. Durable image objects live in a dedicated media directory under the configured AstrBot data root, NEVER the temporary directory or its age-based cleanup policy. Provider adapters continue to receive their existing image content shape after request-time materialization. No extra image-caption model is introduced. - -Alternative rejected: keep base64 in history and only trim it during compaction. This reduces some requests but leaves database growth, repeated JSON parsing, and compaction-record duplication. - -### 2. Keep backward-compatible dual readers - -The history reader accepts both the new reference form and existing inline `data:`/base64 forms. The writer emits references only for newly prepared images. Conversion of existing records is an explicit repair/migration operation with backup, dry-run statistics, hash verification, and rollback by restoring the original history. - -Alternative rejected: automatic in-place migration on first read. It makes a read destructive and risks losing recoverability when the media directory is unavailable. - -### 3. Make encoded bytes a first-class limit - -Image preparation uses independent limits for maximum dimensions and maximum encoded payload bytes. Compute base64 length as `4 * ceil(encoded_file_bytes / 3)` without allocating a base64 copy for every candidate. Retain only the current best candidate. Preserve already-compliant original bytes; otherwise try the original supported format and JPEG quality steps, select a valid candidate without enlarging compliant source content, and encode base64 only at the provider boundary. Preserve transparency when required. If no candidate fits, return a readable size error rather than silently returning the original oversized image. - -Proposed product defaults: preserve existing dimension and quality settings, add `image_compress_options.max_encoded_bytes = 4194304` (4 MiB of base64 per image). This is an AstrBot preparation budget, not a claim about vendor limits. A documented or configured smaller provider limit takes precedence. A known aggregate request limit is checked against the complete serialized request, independently of this per-image budget. Never reinterpret base64 bytes as text tokens. No universal aggregate vendor limit is invented. - -For CUA, preserve the oriented original pixel dimensions and coordinates; try encoding changes without resizing and return a readable error if the budget cannot be met. Do not use an enormous pixel limit as a substitute for an explicit preserve-dimensions flag. Animation follows the selected baseline's frame-selection behavior; an older branch that skips animations is not silently treated as the current montage implementation. Do not keep all decoded animation frames resident simultaneously. - -Alternative rejected: use only pixel dimensions or only source-file bytes. Neither predicts the final JSON payload reliably. - -### 4. Materialize only selected references - -Context selection happens before resolving media references. Images not selected for either the main request or a separate image-capable summary request are not opened, decoded, or encoded. The summary request and the subsequent main request have separate measured lifetimes. Selected references are materialized into a request-local structure and released after the request; the history model never receives the resulting data URI. In-flight cancellation must not delete a file still used by a non-cancelled image worker; cleanup runs when that worker actually exits. - -For a fixed provider configuration, prepared historical bytes are immutable. Do not re-encode the history when a new image arrives, tune old JPEG quality to the remaining request space, or substitute captions to achieve memory targets. Ordinary summary/truncation semantics are unchanged. A known oversized aggregate request fails clearly; this project does not add image eviction or a new compression trigger. B-only lazy loading MUST send byte-identical image content to the baseline. - -Alternative rejected: materialize all references first and let the context manager remove messages afterward. That preserves the current memory spike and defeats lazy loading. - -### 5. Treat all sources as adapters to one preparation function - -Platform attachments, quoted content, plugin/MCP results, file tools, and CUA screenshots normalize to a common input descriptor. The descriptor carries source ownership and cleanup responsibility. The preparation function owns format detection, dimension checks, byte checks, temporary-file cleanup, and diagnostic metadata. - -### 6. Use an external-process ablation harness - -Each measurement case runs in a fresh child process. The harness records RSS at high frequency, Python allocation snapshots separately, request JSON size, image byte sizes, media object counts, and temporary files. Cold-start, warm-start, single-image, multi-image, long-history, concurrent, restart, and missing-media cases are separate workloads. - -The implementation is accepted only when the optimized path preserves image count/order and provider-visible content for the active window. Memory reduction caused by silently dropping images is a failed experiment, not a success. - -## Risks / Trade-offs - -- [Risk] Media references can outlive their files. → Keep content hashes and metadata, report a bounded placeholder on misses, and add cleanup/reconciliation diagnostics. -- [Risk] Existing plugins may assume persisted `image_url` always contains a data URI. → Keep the dual reader and resolve references at the provider boundary; add compatibility tests for plugin/tool history. -- [Risk] Extra disk I/O can increase latency. → Resolve only selected images, cache prepared media by content hash within one request, and measure cold/warm latency separately. -- [Risk] Image conversion can use native memory invisible to `tracemalloc`. → Use an external RSS monitor and report both RSS and Python allocations. -- [Risk] A provider may have a smaller limit than the internal default. → Apply the effective provider limit during request preparation and preserve readable provider-specific errors. - -## Migration Plan - -1. Ship dual reading before new reference writing; factor switches belong in the experiment harness, not a new public rollout switch. -2. Run the ablation suite and compatibility tests before enabling reference writing by default. -3. Provide a dry-run migration report for old inline histories. -4. Migrate selected conversations only after media files and hashes are verified. -5. Before downgrading to a reader without reference support, stop writes and export/re-inline all reference-bearing histories into a verified backup, including conversations created after deployment. Merely disabling new writes does not make those records readable by old versions. Retain the media store until rollback verification completes. - -## 兼容与资源生命周期约束 - -- 当前基线工作区提交为 `1a2492a09f8722a9abd813f6a65fb687d53c88c9`,位于此前功能分支,不等同于 issue 的 v4.28.1。实现前固定目标基线及依赖,另建 v4.28.1 只读复现环境;不能拿两个版本的路径混合计算收益。 -- 新引用采用内部版本化图片部件,保留 image detail;对外 ProviderRequest 输入继续接受原来的路径、URL、data URI。新类型不得直接发送给远端 Provider。 -- Plugin/第三方 Provider 的兼容出口必须解析引用后再调用现有接口;内置 Provider 和摘要 Provider 都纳入测试。给插件的兼容视图可能仍需分配 base64,这部分成本如实报告,不能声称所有消费者都已懒加载。 -- 媒体文件采用内容哈希去重,先在同文件系统原子写入并核验,再提交消息引用。写入失败保持原会话不变;崩溃留下的未引用文件由维护命令处理。 -- 媒体读取通过内部 ID 查询并限制路径;WebUI 使用带会话权限校验的图片端点,不公开数据目录或仅凭哈希授权。会话列表/详情只返回引用和预览地址,不批量还原 base64。 -- 复用现有附件存储能力前核对其 WebChat 删除逻辑,禁止把 Agent 媒体直接挂入会被其他界面独立删除的生命周期。引用对象与消息内容绑定,跨会话去重不得带来跨会话读权限。 -- 首版不做运行中自动垃圾回收。显式维护命令在停止写入后扫描全部保留会话,输出 dry-run,隔离未引用对象;读写中断、解析失败时停止删除。移除或压缩一个会话不能删除仍被其他会话引用的文件。 -- 旧 inline 图片在普通读取时不转码、不改写;显式迁移只搬运原字节并核验哈希,不在迁移中夹带画质调整。历史包与媒体目录共同备份;导出保持自包含或附带媒体清单,导入必须校验引用。 -- 本地资源异常不应被笼统吞掉再走同一大请求;测试检查 MemoryError 保留异常类型并终止本次运行。413 与资源耗尽分开记录,不假设每个回退 Provider 的上限都相同,不引入自动多轮重压缩循环。 -- 这不是对话事件日志重构;当前正常上下文压缩及保存行为保持不变。媒体表示迁移与“保留所有原始对话”的产品需求不得混为一谈。 - -## 详细内存消融实验 - -### 因素与对照 - -三个开关仅供实验调用真实生产函数,不提供给用户。完整八组为 `000 / A00 / 0B0 / 00C / AB0 / A0C / 0BC / ABC`。 - -| 因素 | 唯一变化 | 必须保持不变 | -| --- | --- | --- | -| A 统一入口 | 所有来源进入同一准备链 | 使用原压缩算法和参数 | -| B 懒加载 | 同一准备结果外置存储,选窗后加载 | 图片字节、MIME、顺序、detail、消息位置 | -| C 压缩优化 | 仅在基线已有入口改变压缩算法 | 不顺便补齐工具/插件入口 | - -八组的零开关路径必须与固定基线的输入输出一致;用独立基线 checkout 的相同函数校验,不能把一个近似模拟器叫基线。另加纯文本对照,但不把它算作三因素收益。 - -### 固定样本 - -五类图片各包含普通与压力样本,每类三个固定种子,至少 30 个实例。文件在测量进程外生成,记录 SHA256、编码格式、尺寸、方向、透明度、帧数和字节数。 - -| 类别 | 普通 | 压力 | -| --- | --- | --- | -| PNG | 透明图/色块 | 1280×1280 高熵 RGBA,像素合规但字节超限 | -| JPEG | 常见照片、已压缩小文件 | 4000×3000 高细节、EXIF 旋转 | -| WebP | 静态有损 | 无损透明,动画分支另做覆盖 | -| GIF | 少帧动画 | 多帧动画,检查帧抽取/拼图峰值 | -| 工具截图 | 1920×1080 终端 | 4K 小字 UI 与坐标标记 | - -来源测试覆盖本地文件、HTTP、data URI、base64 URI,以及用户附件、引用、插件/MCP、FileRead、CUA。HTTP 使用本地服务器;这些入口调用现有适配代码,不能仅给同一函数贴不同来源名称。完整交叉用于单图工作负载,其余长会话用五种格式均衡混合。 - -### 工作负载 - -1. 单图单轮、同轮八张不同图片;分别统计每个阶段。 -2. 50 轮每轮一张不同图片,然后 20 轮纯文本;与重复同一张图片的独立测试区分去重收益。 -3. 10/50/100 张历史图;分别测全在窗口、按轮次裁剪、摘要成功和摘要失败四种状态。记录实际触发次数,不假设默认配置一定触发。 -4. 摘要模型分别支持和不支持图片。固定摘要返回值及用量,避免随机文本干扰;摘要图像输入遵循实际模态规则。 -5. 四会话并发、保存后重启、WebUI 只看详情/展开单张预览、显式导出;WebUI 与模型请求分别统计。 -6. 错误实验:缺失媒体、坏图片、取消、超时、写盘失败、5 MiB 聚合请求拒绝、独立单图限制、不同限制的回退 Provider。 -7. 静态图片循环固定窗口请求 200 次,记录 1/10/25/50/100/200 次,检查存活对象、临时文件、打开句柄及运行后残留;不强制 GC 的正常数据与诊断 GC 分开。 - -### 测量与防止实验污染 - -- 分阶段打点:读取数据库 → JSON 解析/消息构造 → 选窗 → 读取媒体 → 解码/缩放/编码 → base64 → Provider 组装 → HTTP 序列化/发送 → 保存 → 清理。另标记摘要请求阶段。 -- 记录原文件与输出文件字节、base64 长度、完整 HTTP JSON 字节、读取/解码/转码/base64 次数及字节量、阶段耗时、临时文件和持久化字节。磁盘数据不等价于内存收益。 -- 外部监测进程每 10 ms 采集被测进程及子进程 RSS;记录 OS 高水位。Windows 另采私有内存与工作集。两个平台单独比较,不直接比较绝对值。 -- 标准轮关闭 tracemalloc;诊断轮单独启用 tracemalloc,Linux 关键图另外追踪原生分配。先前约 24 MB WebP 结果缺乏初始化隔离,不能用作验收基线。 -- 每个实验单元十次独立进程运行,冷启动/预热三次后的热运行分别记录。按种子配对,随机交错八组,保存完整运行顺序。固定机器、解释器、依赖和处理并发数。 -- 本地模型服务返回固定文本、固定用量和固定摘要;真实客户端/SDK 必须参与最终请求组装。服务端流式计数,不保留所有历史请求体;结构校验另跑,不污染测量峰值。 -- 监测器不拷贝 base64、不输出原图、会话正文或密钥。样本生成、结果分析和监测服务不得计入被测进程。 -- 单次 120 秒或被测进程总内存达到 `min(2 GiB, 启动时可用内存的 25%)` 时停止并记录超限。崩溃、超时、OOM 数据不得丢弃或用零内存填充。 - -### 分析和验收 - -- 单因素效果比较 A/B/C 与基线,边际效果比较 ABC 与 BC/AC/AB;报告配对中位差、范围与原始样本,不把三个百分比相加。十次重复不宣称精确的 P99。 -- B-only 校验图片哈希、顺序、detail 和固定 Provider 的历史前缀;预期最终请求大小不变。请求时仍需装载整个有效图片窗口,必须披露这个下限。 -- C 校验方向、透明、截图文字和坐标;已合规原字节应原样保留,不能选用无必要膨胀的输出。超限失败不算“成功压缩”或内存收益。 -- CUA 的尺寸和坐标必须一致;失败不能伪装成空图成功。GIF 的既有抽帧语义必须一致。 -- B 在 50 张历史图的读取到选窗阶段,以相同消息窗口为前提,峰值增量降低至少 50%;ABC 在高熵长会话的成功可比场景中,整体峰值增量降低至少 30%。这是验收目标,不是已证明结果。 -- 小图正常场景峰值增量及中位耗时回退不超过 10%;图片质量或稳定历史前缀失败时,即使内存达标也不通过。 -- 持续请求后存活图片对象、文件句柄、临时文件不能按轮数线性增长。RSS 分配器保留不直接判定泄漏;结合对象与原生分配数据解释。 -- 实测平台最低覆盖 Linux/WSL 与原生 Windows,macOS 做功能回归。所有指标原始 JSONL/CSV、样本清单和曲线放在测试产物目录,不创建仓库 SUMMARY 报告。 -- 真实智谱只用于可选小样本画质/协议验证,与消融数据隔离。默认不发公网请求;密钥通过安全环境注入,绝不写入本规划或测试产物。 diff --git a/openspec/changes/optimize-image-memory-lifecycle/proposal.md b/openspec/changes/optimize-image-memory-lifecycle/proposal.md deleted file mode 100644 index 8d6ffbcaa3..0000000000 --- a/openspec/changes/optimize-image-memory-lifecycle/proposal.md +++ /dev/null @@ -1,28 +0,0 @@ -## Why - -AstrBot currently prepares images through several partially independent paths. Some inputs are resized only by pixel dimensions, while the final base64 payload can still exceed provider limits and consume several copies of the image in memory. Conversation history can also retain complete inline image data, so every later request pays the deserialization and request-construction cost again. Issue #10089 demonstrates request-size failures, while #10092 demonstrates process-level memory exhaustion. - -## What Changes - -- Add one image preparation contract for user attachments, quoted images, plugins, MCP/tool results, file reading, and screenshots. -- Enforce both dimension limits and a configurable final encoded-payload byte limit; reject or degrade images that cannot satisfy the limit instead of silently sending oversized data. -- Store historical images as durable media references and metadata rather than embedding complete base64 data in every conversation message. -- Resolve historical image references only when the active provider request needs them, with bounded loading and cleanup. -- Keep the current conversation semantics: images in the active context remain available to the model; this change does not silently remove images merely to improve memory numbers. -- Add memory-ablation tests covering preparation, history loading, request assembly, persistence, and repeated conversations. -- Preserve backward compatibility by reading existing inline data URIs and provide an explicit migration/repair path instead of rewriting them implicitly. - -## Capabilities - -### New Capabilities - -- `image-memory-lifecycle`: bounded image preparation, durable media references, lazy resolution, and memory-safe request assembly. -- `agent-context-image-budget`: active requests validate image byte size independently from token accounting and resolve references without changing image visibility. -- `conversation-history-media`: persisted messages may use media references and must remain readable across restart while preserving existing inline-image histories. - -## Impact - -- Affects media resolution and compression, provider request assembly, agent context processing, conversation persistence, tool/plugin image paths, and related tests. -- Adds a local media storage/index lifecycle with cleanup and missing-media handling. -- Adds configuration for final encoded image size and bounded media storage behavior; existing pixel/quality settings remain compatible. -- Provider payloads remain provider-specific; only the internal history representation and preparation boundary change. diff --git a/openspec/changes/optimize-image-memory-lifecycle/specs/agent-context-image-budget/spec.md b/openspec/changes/optimize-image-memory-lifecycle/specs/agent-context-image-budget/spec.md deleted file mode 100644 index 0e14546564..0000000000 --- a/openspec/changes/optimize-image-memory-lifecycle/specs/agent-context-image-budget/spec.md +++ /dev/null @@ -1,40 +0,0 @@ -## Purpose - -Keeps the active model context bounded by treating image payload bytes as a real request cost and by resolving only image data required by the active provider request. - -## ADDED Requirements - -### Requirement: Active context SHALL use a bounded image budget -The system SHALL validate encoded image bytes and known provider request-size limits independently from text token estimates. It MUST NOT introduce additional image eviction, replace available images with summaries, or re-encode historical images on every turn to fit a changing budget. - -#### Scenario: Context contains many images -- **WHEN** the active context contains images whose encoded payloads exceed the image budget -- **THEN** the system reports a readable size-limit error before sending a known-oversized request, without silently dropping images or changing historical image bytes - -### Requirement: Historical media SHALL be resolved on demand -The system SHALL keep historical image content as a resolvable media reference and SHALL materialize its bytes only when the image is selected for the active provider request. - -#### Scenario: Old image is outside the active context -- **WHEN** a historical image is outside the selected context window and is not input to the summarization request -- **THEN** the system does not read, decode, or base64-encode that image for the request - -### Requirement: Summarization SHALL retain its existing image input semantics -The system SHALL resolve selected image references for an image-capable summarization provider just as it does for the main provider. Images summarized out of the main window may still need loading for that separate summary request; this cost SHALL be measured separately. - -#### Scenario: Old image participates in a summary -- **WHEN** the existing context compression policy passes an old image to an image-capable summary provider -- **THEN** that provider receives the image rather than an unresolved reference, while the subsequent main request contains only its selected context - -### Requirement: Lazy loading SHALL preserve provider-visible image bytes -The system SHALL preserve image bytes, MIME type, detail, order, and message placement for a fixed prepared image across persistence, restart, and subsequent requests to the same provider configuration. - -#### Scenario: A later turn loads a persisted image -- **WHEN** a later text message causes an existing image reference to be materialized -- **THEN** the image content is byte-identical to its earlier prepared representation and the stable historical request prefix remains unchanged - -### Requirement: Existing inline images SHALL remain readable -The system SHALL continue to read existing persisted inline image data and SHALL apply the active image budget when such data is included in a request. - -#### Scenario: Existing conversation contains a data URI -- **WHEN** a conversation created before media references is loaded -- **THEN** the conversation remains readable and its inline image is handled by the same bounded request policy diff --git a/openspec/changes/optimize-image-memory-lifecycle/specs/conversation-history-media/spec.md b/openspec/changes/optimize-image-memory-lifecycle/specs/conversation-history-media/spec.md deleted file mode 100644 index 23201fa078..0000000000 --- a/openspec/changes/optimize-image-memory-lifecycle/specs/conversation-history-media/spec.md +++ /dev/null @@ -1,37 +0,0 @@ -## Purpose - -Separates durable conversation facts from large media payloads so that conversations remain restartable and image-capable without duplicating complete base64 data in every historical message. - -## ADDED Requirements - -### Requirement: New persisted image content SHALL use media references -The system SHALL persist a stable media reference with MIME type and image metadata for newly saved model-visible images instead of embedding the complete base64 payload in the conversation record. - -#### Scenario: Model-visible image is saved -- **WHEN** a request containing a newly prepared image is saved to conversation history -- **THEN** the history stores a media reference and sufficient metadata to resolve the image later - -### Requirement: Media references SHALL survive restart -The system SHALL resolve persisted media references after restart and SHALL report a readable missing-media result when the referenced media is unavailable. - -#### Scenario: Referenced media file is missing -- **WHEN** a conversation contains a reference whose media object no longer exists -- **THEN** loading the conversation does not crash, the UI reports unavailable media, and the model receives an explicit bounded missing-image placeholder; this is reported as a recovery condition, not a successful memory optimization - -### Requirement: Media lifetime SHALL be independent of temporary cleanup -Referenced media SHALL be durable outside the temporary-file cleanup domain. The system SHALL authorize image reads through conversation access, preserve shared media until no retained history references it, and never interpret a client-controlled media reference as an arbitrary local path. - -#### Scenario: Temporary files are cleaned after restart -- **WHEN** temporary media cleanup runs while a saved conversation still references an image -- **THEN** that saved image remains readable - -#### Scenario: Another conversation requests an unauthorized image -- **WHEN** a caller supplies an image identifier without access to its owning conversation -- **THEN** the request is rejected without exposing the image or its filesystem path - -### Requirement: History migration SHALL be explicit and reversible -The system SHALL preserve existing inline-image histories and SHALL provide an explicit migration or repair operation before replacing inline payloads with media references. - -#### Scenario: Existing history has inline base64 -- **WHEN** an old conversation is opened without migration -- **THEN** it remains readable and no destructive rewrite occurs automatically diff --git a/openspec/changes/optimize-image-memory-lifecycle/specs/image-memory-lifecycle/spec.md b/openspec/changes/optimize-image-memory-lifecycle/specs/image-memory-lifecycle/spec.md deleted file mode 100644 index 871a21006e..0000000000 --- a/openspec/changes/optimize-image-memory-lifecycle/specs/image-memory-lifecycle/spec.md +++ /dev/null @@ -1,26 +0,0 @@ -## Purpose - -Provides one bounded lifecycle for images from users, tools, plugins, and external media so that model requests remain within provider limits without multiplying large image payloads in memory. - -## ADDED Requirements - -### Requirement: All model-bound images use one bounded preparation contract -The system SHALL apply the same preparation contract to user attachments, quoted images, plugin and MCP results, file-reading images, and computer-use screenshots before they become model image content. - -#### Scenario: Source-specific image enters a model request -- **WHEN** an image is supplied by any supported source -- **THEN** the system applies the same dimension, format, encoded-byte, cleanup, and failure rules before constructing provider content - -### Requirement: Image preparation SHALL enforce final encoded size -The system SHALL enforce a configurable upper bound on the final encoded image payload, in addition to pixel dimensions, and SHALL NOT silently send an image that exceeds the bound after preparation. - -#### Scenario: Pixel-compliant image exceeds the byte limit -- **WHEN** an image is within the configured dimensions but its encoded payload exceeds the configured byte limit -- **THEN** the system re-encodes or resizes it until it fits, or returns a readable image-too-large failure without sending the oversized payload - -### Requirement: Preparation SHALL release temporary resources -The system SHALL clean resolver-owned temporary files and release image-processing resources after request construction succeeds, fails, is cancelled, or times out. - -#### Scenario: Image preparation is cancelled -- **WHEN** request processing is cancelled during download, decode, resize, or encoding -- **THEN** temporary files and owned buffers are released without affecting unrelated media diff --git a/openspec/changes/optimize-image-memory-lifecycle/tasks.md b/openspec/changes/optimize-image-memory-lifecycle/tasks.md deleted file mode 100644 index 22982703fa..0000000000 --- a/openspec/changes/optimize-image-memory-lifecycle/tasks.md +++ /dev/null @@ -1,36 +0,0 @@ -## 1. Baseline and measurement harness - -- [x] 1.1 Freeze the implementation and dependency baseline, create a separate read-only v4.28.1 reproduction environment for issues #10089 and #10092, and verify both environments report their commit and dependency versions. -- [x] 1.2 Add an external-process image benchmark harness that records stage timings, RSS high-water marks, Python allocation diagnostics, child-process memory, request JSON bytes, media bytes, temporary files, and failures without printing image data or secrets. -- [x] 1.3 Generate deterministic PNG, JPEG, WebP, GIF, and screenshot fixtures with ordinary and high-entropy cases, record hashes and metadata, and verify the fixture generator is outside the measured child process. -- [ ] 1.4 Implement the eight-factor ablation runner for unified-source preparation, lazy history loading, and compression optimization; verify paired seeds, cold/warm runs, ten repetitions, and raw JSONL/CSV output. - -## 2. Unified preparation and encoded-byte limits - -- [x] 2.1 Define the internal image preparation descriptor and one preparation entry point for platform attachments, quoted images, plugins/MCP, file tools, and CUA screenshots; verify every listed source reaches it with source ownership and cleanup metadata. -- [x] 2.2 Extend image preparation with an encoded payload byte limit independent from pixel dimensions; verify a pixel-compliant high-entropy PNG is re-encoded or rejected instead of passed through. -- [x] 2.3 Implement candidate-size measurement without retaining all candidate bytes and preserve transparency, EXIF orientation, CUA dimensions, animation behavior, and already-compliant bytes; verify format-specific fixtures and coordinate/visual checks. -- [x] 2.4 Add typed, readable image-size and resource errors and ensure oversize input is not silently returned to the Provider or retried through unrelated fallback providers; verify #10089's request-size scenario. -- [x] 2.5 Add cancellation, timeout, decode failure, and write failure cleanup tests; verify temporary files, image-library resources, and worker tasks are released on every path. - -## 3. Durable media references and lazy materialization - -- [x] 3.1 Add versioned durable media objects under the configured data directory with content-hash deduplication, atomic write/verify, MIME/dimension/byte/detail metadata, and conversation-scoped access checks; verify partial writes never create usable references. -- [x] 3.2 Add a persisted image-reference representation and dual history reader for new references plus existing inline data URIs/base64; verify old histories remain readable and new histories do not embed complete base64. -- [x] 3.3 Resolve selected references only after context selection and materialize them into a request-local provider view; verify out-of-window images are not opened/decoded/encoded and B-only provider-visible bytes, order, detail, and placement match baseline. -- [x] 3.4 Keep durable media outside temporary cleanup and add missing-media, restart, unauthorized-reference, session-delete, and cross-session shared-media tests; verify missing media produces a bounded recovery result without exposing paths. -- [x] 3.5 Add explicit dry-run migration/export/repair behavior for old inline histories with backup and hash verification; verify ordinary reads never rewrite history and rollback remains possible. - -## 4. Context and persistence integration - -- [x] 4.1 Integrate image byte validation with the active request and summarization-provider paths without adding image eviction or changing existing truncation/summary semantics; verify image-capable summaries receive selected images and the main request receives its selected context. -- [x] 4.2 Ensure history saving persists references rather than request-time data URIs while provider adapters continue receiving their existing wire shapes; verify plugin, tool, built-in Provider, WebUI detail, and export compatibility. -- [x] 4.3 Separate 413 request-size, image-too-large, missing-media, and MemoryError handling; verify MemoryError preserves its type and does not trigger repeated oversized fallback requests. - -## 5. Verification and rollout - -- [ ] 5.1 Run the full matrix across single image, eight images, 50-round history, 200 fixed-window requests, four-session concurrency, restart, WebUI preview, and error workloads; verify all metrics and raw failures are retained. -- [ ] 5.2 Compare A/B/C main and marginal effects with image-count/order/content invariants; verify memory gains are not caused by silently omitting active images. -- [ ] 5.3 Validate Linux/WSL and native Windows memory behavior, run macOS functional regression, and verify normal small-image latency/peak-memory regression stays within the documented target. -- [x] 5.4 Run `ruff format --check .`, `ruff check .`, focused image/context/history tests, and the relevant full test suites; verify no source comments/logs violate repository language rules. -- [ ] 5.5 Review migration backup/rollback and media cleanup behavior, document measured results in test artifacts outside repository SUMMARY files, and only then enable new reference writes by default. From 4eb2233668a90d02600af977a5eb9d5694678d50 Mon Sep 17 00:00:00 2001 From: zenfun Date: Thu, 17 Sep 2026 18:27:44 +0800 Subject: [PATCH 6/7] chore: remove image memory experiment tooling --- .github/workflows/image-memory-validation.yml | 215 -------- scripts/image_memory_bench/ablation_runner.py | 66 --- .../image_memory_bench/baseline_manifest.py | 54 -- .../image_memory_bench/generate_fixtures.py | 127 ----- .../image_memory_bench/history_benchmark.py | 461 ------------------ .../image_memory_bench/lifecycle_workloads.py | 371 -------------- scripts/image_memory_bench/measure.py | 399 --------------- scripts/image_memory_bench/migrate_history.py | 99 ---- .../prepare_lifecycle_fixture.py | 131 ----- scripts/manage_image_history.py | 434 ----------------- tests/unit/test_image_lifecycle_workloads.py | 36 -- tests/unit/test_image_memory_benchmark.py | 59 --- tests/unit/test_manage_image_history.py | 354 -------------- 13 files changed, 2806 deletions(-) delete mode 100644 .github/workflows/image-memory-validation.yml delete mode 100644 scripts/image_memory_bench/ablation_runner.py delete mode 100644 scripts/image_memory_bench/baseline_manifest.py delete mode 100644 scripts/image_memory_bench/generate_fixtures.py delete mode 100644 scripts/image_memory_bench/history_benchmark.py delete mode 100644 scripts/image_memory_bench/lifecycle_workloads.py delete mode 100644 scripts/image_memory_bench/measure.py delete mode 100644 scripts/image_memory_bench/migrate_history.py delete mode 100644 scripts/image_memory_bench/prepare_lifecycle_fixture.py delete mode 100644 scripts/manage_image_history.py delete mode 100644 tests/unit/test_image_lifecycle_workloads.py delete mode 100644 tests/unit/test_image_memory_benchmark.py delete mode 100644 tests/unit/test_manage_image_history.py diff --git a/.github/workflows/image-memory-validation.yml b/.github/workflows/image-memory-validation.yml deleted file mode 100644 index a1ff61278c..0000000000 --- a/.github/workflows/image-memory-validation.yml +++ /dev/null @@ -1,215 +0,0 @@ -name: Image Memory Lifecycle Validation - -on: - push: - branches: - - feat/optimize-image-memory-lifecycle - workflow_dispatch: - -jobs: - validate: - name: Image and memory validation (${{ matrix.os }}) - runs-on: ${{ matrix.os }} - timeout-minutes: 90 - strategy: - fail-fast: false - matrix: - os: [ubuntu-latest, windows-latest, macos-latest] - - defaults: - run: - shell: bash - - steps: - - name: Checkout - uses: actions/checkout@v7 - - - name: Set up Python - uses: actions/setup-python@v7 - with: - python-version: '3.12' - - - name: Install uv - run: | - python -m pip install --upgrade pip - python -m pip install uv - - - name: Install dependencies - run: uv sync - - - name: Run repository ruff checks - run: | - uv run ruff format --check . - uv run ruff check . - - - name: Run focused tests - run: | - mkdir -p "$RUNNER_TEMP/image-memory-validation/results" - uv run pytest -q \ - tests/unit/test_image_preparation_budget.py \ - tests/unit/test_image_media_store.py \ - tests/test_media_utils.py \ - tests/test_computer_fs_tools.py \ - tests/test_tool_loop_agent_runner.py \ - tests/agent/test_context_manager.py \ - tests/unit/test_astr_main_agent.py \ - tests/unit/test_context_image_budget.py \ - tests/unit/test_image_source_preparation.py \ - tests/unit/test_image_memory_benchmark.py \ - tests/unit/test_image_lifecycle_workloads.py \ - tests/unit/test_image_screenshot_encoding.py \ - tests/unit/test_image_source_integration.py \ - tests/unit/test_provider_image_references.py \ - tests/test_agent_runner_media_resolver.py \ - tests/unit/test_provider_request_retry.py \ - tests/unit/test_manage_image_history.py \ - tests/unit/test_conversation_media_api.py \ - tests/test_openai_source.py \ - tests/test_openai_responses_source.py \ - tests/test_gemini_source.py \ - tests/test_anthropic_source.py \ - tests/test_anthropic_kimi_code_provider.py \ - tests/test_deerflow_agent_runner.py \ - tests/unit/test_cron_context_compression.py \ - tests/unit/test_group_chat_context_wiring.py \ - tests/test_platform_image_format_preservation.py \ - tests/test_webchat_queue_lifecycle.py \ - tests/test_webchat_upload_image_format.py \ - --junitxml="$RUNNER_TEMP/image-memory-validation/results/tests.xml" - - - name: Generate reproducible fixtures outside measured processes - run: | - mkdir -p "$RUNNER_TEMP/image-memory-validation/fixtures" - uv run python scripts/image_memory_bench/generate_fixtures.py \ - "$RUNNER_TEMP/image-memory-validation/fixtures" \ - --seeds 7 19 43 - - - name: Record versions and fixture manifest - run: | - mkdir -p "$RUNNER_TEMP/image-memory-validation/results" - uv run python scripts/image_memory_bench/baseline_manifest.py \ - "$RUNNER_TEMP/image-memory-validation/results/manifest.json" \ - --fixtures "$RUNNER_TEMP/image-memory-validation/fixtures" - - - name: Run single-image memory workloads for every paired seed - run: | - set +e - result="$RUNNER_TEMP/image-memory-validation/results/measurements.jsonl" - failures=0 - for seed in 7 19 43; do - while IFS= read -r image; do - uv run python scripts/image_memory_bench/measure.py \ - "$image" "$result" \ - --repo "$GITHUB_WORKSPACE" \ - --csv "$RUNNER_TEMP/image-memory-validation/results/measurements.csv" \ - --workload "$seed-$(basename "$image")" - code=$? - if [ "$code" -ne 0 ]; then - failures=$((failures + 1)) - fi - done < <(find "$RUNNER_TEMP/image-memory-validation/fixtures/$seed" \ - -type f \( -name '*ordinary*' -o -name '*stress*' \) | sort) - done - printf '{"measurement_failures":%s}\n' "$failures" \ - > "$RUNNER_TEMP/image-memory-validation/results/measurement-summary.json" - - - name: Run Python allocation diagnostic workload - run: | - uv run python scripts/image_memory_bench/measure.py \ - "$RUNNER_TEMP/image-memory-validation/fixtures/7/png-stress.png" \ - "$RUNNER_TEMP/image-memory-validation/results/trace-python.jsonl" \ - --repo "$GITHUB_WORKSPACE" \ - --trace-python \ - --csv "$RUNNER_TEMP/image-memory-validation/results/trace-python.csv" \ - --workload trace-python-png-stress - - - name: Run paired history loading workloads - run: | - mkdir -p "$RUNNER_TEMP/image-memory-validation/results/history" - for seed in 7 19 43; do - uv run python scripts/image_memory_bench/history_benchmark.py \ - "$RUNNER_TEMP/image-memory-validation/fixtures/$seed/manifest.json" \ - "$RUNNER_TEMP/image-memory-validation/results/history/$seed.jsonl" \ - --repo "$GITHUB_WORKSPACE" \ - --images 50 \ - --keep-turns 25 \ - --repeats 10 \ - --cold-repeats 10 \ - --warm-repeats 10 \ - --distinct \ - --csv "$RUNNER_TEMP/image-memory-validation/results/history/$seed.csv" \ - --rss-jsonl "$RUNNER_TEMP/image-memory-validation/results/history/$seed.rss.jsonl" - done - - - name: Run fixed-window concurrent lifecycle and restart workload - run: | - lifecycle_root="$RUNNER_TEMP/image-memory-validation/lifecycle" - mkdir -p "$lifecycle_root" - uv run python scripts/image_memory_bench/prepare_lifecycle_fixture.py \ - "$RUNNER_TEMP/image-memory-validation/fixtures/7/manifest.json" \ - "$lifecycle_root/input" \ - --session-count 4 \ - --history-turns 50 - lifecycle_manifest="$lifecycle_root/input/lifecycle-manifest.json" - database=$(uv run python -c \ - 'import json,sys; print(json.load(open(sys.argv[1]))["database"])' \ - "$lifecycle_manifest") - media_root=$(uv run python -c \ - 'import json,sys; print(json.load(open(sys.argv[1]))["media_root"])' \ - "$lifecycle_manifest") - conversation_ids=$(uv run python -c \ - 'import json,sys; print(" ".join(json.load(open(sys.argv[1]))["conversation_ids"]))' \ - "$lifecycle_manifest") - command=( - uv run python scripts/image_memory_bench/lifecycle_workloads.py - "$lifecycle_root/result.json" - --repo "$GITHUB_WORKSPACE" - --database "$database" - --media-root "$media_root" - --session-count 4 - --request-count 200 - --window-turns 25 - --rss-jsonl "$lifecycle_root/result.rss.jsonl" - ) - for conversation_id in $conversation_ids; do - command+=(--conversation-id "$conversation_id") - done - "${command[@]}" - - - name: Record ablation contract status without fabricating factors - run: | - uv run python scripts/image_memory_bench/ablation_runner.py \ - "$RUNNER_TEMP/image-memory-validation/results/ablation-plan.json" \ - --baseline-repo "$GITHUB_WORKSPACE" \ - --candidate-repo "$GITHUB_WORKSPACE" - - - name: Upload raw validation artifacts - if: always() - uses: actions/upload-artifact@v4 - with: - name: image-memory-validation-${{ matrix.os }} - path: ${{ runner.temp }}/image-memory-validation - if-no-files-found: error - - dashboard: - name: Dashboard typecheck - runs-on: ubuntu-latest - timeout-minutes: 20 - defaults: - run: - working-directory: dashboard - steps: - - name: Checkout - uses: actions/checkout@v7 - - name: Set up Node - uses: actions/setup-node@v7 - with: - node-version: 22 - cache: pnpm - cache-dependency-path: dashboard/pnpm-lock.yaml - - name: Enable pnpm - run: corepack enable - - name: Install dashboard dependencies - run: pnpm install --frozen-lockfile - - name: Run dashboard typecheck - run: pnpm typecheck diff --git a/scripts/image_memory_bench/ablation_runner.py b/scripts/image_memory_bench/ablation_runner.py deleted file mode 100644 index 9c8b08abb9..0000000000 --- a/scripts/image_memory_bench/ablation_runner.py +++ /dev/null @@ -1,66 +0,0 @@ -"""Run the real 2^3 image experiment only when production factor flips exist. - -The current checkout does not expose independent A, B, and C experiment -switches. This runner therefore refuses to fabricate a matrix and emits an -auditable blocked manifest until each factor is supplied by a real checkout. -""" - -from __future__ import annotations - -import argparse -import itertools -import json -import subprocess -from pathlib import Path - - -def git_identity(repo: Path) -> dict[str, str]: - """Return commit and dirty diff identity for a checkout.""" - commit = subprocess.check_output( - ["git", "rev-parse", "HEAD"], cwd=repo, text=True - ).strip() - diff = subprocess.check_output(["git", "diff", "--binary"], cwd=repo) - import hashlib - - return { - "commit": commit, - "working_tree_diff_sha256": hashlib.sha256(diff).hexdigest(), - } - - -def main() -> int: - """Validate factor provenance and write a blocked or runnable plan.""" - parser = argparse.ArgumentParser() - parser.add_argument("output", type=Path) - parser.add_argument("--baseline-repo", type=Path, required=True) - parser.add_argument("--candidate-repo", type=Path, required=True) - parser.add_argument("--factor-entrypoint", type=Path) - args = parser.parse_args() - factors = { - "A": "unavailable: no independent production source-preparation entrypoint supplied", - "B": "unavailable: no independent production storage/lazy-materialization entrypoint supplied", - "C": "unavailable: no independent production compression-algorithm entrypoint supplied", - } - entrypoint_ok = ( - args.factor_entrypoint is not None and args.factor_entrypoint.is_file() - ) - plan = { - "status": "blocked" if not entrypoint_ok else "requires_factor_contract", - "experiment": "real-2^3-factorial", - "factors": factors, - "combinations": ["".join(bits) for bits in itertools.product("0A", repeat=0)], - "baseline": git_identity(args.baseline_repo), - "candidate": git_identity(args.candidate_repo), - "warmup_scope": "sdk_send_only; not warm history", - "reason": "Do not execute or label combinations until A/B/C independently select real production paths.", - } - plan["combinations"] = [ - "".join(bits) for bits in itertools.product("0A", "0B", "0C") - ] - args.output.write_text(json.dumps(plan, indent=2) + "\n", encoding="utf-8") - print(json.dumps({"status": plan["status"], "output": str(args.output)})) - return 0 if plan["status"] == "blocked" else 2 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/scripts/image_memory_bench/baseline_manifest.py b/scripts/image_memory_bench/baseline_manifest.py deleted file mode 100644 index 1710c775cb..0000000000 --- a/scripts/image_memory_bench/baseline_manifest.py +++ /dev/null @@ -1,54 +0,0 @@ -"""Record the exact code and dependency baseline for image experiments.""" - -from __future__ import annotations - -import argparse -import hashlib -import importlib.metadata -import json -import platform -import subprocess -import sys -from pathlib import Path - - -def main() -> int: - """Write a reproducibility manifest without reading application data.""" - parser = argparse.ArgumentParser() - parser.add_argument("output", type=Path) - parser.add_argument("--fixtures", type=Path) - args = parser.parse_args() - commit = subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip() - packages = sorted( - f"{dist.metadata['Name']}=={dist.version}" - for dist in importlib.metadata.distributions() - if dist.metadata.get("Name") - ) - manifest = { - "commit": commit, - "working_tree_diff_sha256": hashlib.sha256( - subprocess.check_output(["git", "diff", "--binary"], cwd=Path.cwd()) - ).hexdigest(), - "python": sys.version, - "platform": platform.platform(), - "packages": packages, - } - if args.fixtures: - files = [] - for path in sorted(args.fixtures.rglob("*")): - if path.is_file() and path.name != "manifest.json": - files.append( - { - "path": str(path.relative_to(args.fixtures)), - "bytes": path.stat().st_size, - "sha256": hashlib.sha256(path.read_bytes()).hexdigest(), - } - ) - manifest["fixtures"] = files - args.output.write_text(json.dumps(manifest, indent=2) + "\n", encoding="utf-8") - print(json.dumps({"commit": commit, "package_count": len(packages)})) - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/scripts/image_memory_bench/generate_fixtures.py b/scripts/image_memory_bench/generate_fixtures.py deleted file mode 100644 index 74c8cc9a3a..0000000000 --- a/scripts/image_memory_bench/generate_fixtures.py +++ /dev/null @@ -1,127 +0,0 @@ -"""Generate seeded fixtures outside the measured application process.""" - -from __future__ import annotations - -import argparse -import hashlib -import json -import random -from pathlib import Path - -from PIL import Image, ImageDraw - - -def generate(output: Path, seed: int) -> list[dict]: - """Create ten inputs for one paired experiment seed. - - Args: - output: New directory for this seed's fixtures. - seed: Stable seed for pixel generation. - - Returns: - On-disk hashes and decoded image metadata. - """ - output.mkdir(parents=True, exist_ok=False) - rng = random.Random(seed) - records = [] - for kind in ("png", "jpeg", "webp", "gif", "screenshot"): - for stress in (False, True): - name = f"{kind}-{'stress' if stress else 'ordinary'}" - options = {} - if kind in ("png", "webp"): - size = (1280, 1280) if stress else (320, 240) - image = ( - Image.frombytes("RGBA", size, rng.randbytes(size[0] * size[1] * 4)) - if stress - else Image.new("RGBA", size, (40, 120, 200, 127)) - ) - image_format = kind.upper() - options = {"lossless": stress} if kind == "webp" else {} - elif kind == "jpeg": - size = (4000, 3000) if stress else (640, 480) - image = Image.frombytes( - "RGB", size, rng.randbytes(size[0] * size[1] * 3) - ) - image_format = "JPEG" - options = {"quality": 95 if stress else 75} - if stress: - exif = image.getexif() - exif[274] = 6 - options["exif"] = exif - elif kind == "gif": - image = Image.new("RGB", (320, 240), (seed % 255, 0, 255)) - image_format = "GIF" - frames = [] - for frame in range(24 if stress else 3): - extra = Image.new("RGB", image.size, (frame * 9, 80, 20)) - ImageDraw.Draw(extra).rectangle( - (frame * 5, 10, frame * 5 + 30, 50), fill="white" - ) - frames.append(extra) - options = { - "save_all": True, - "append_images": frames, - "duration": 80, - "loop": 0, - } - else: - image = Image.new( - "RGB", (3840, 2160) if stress else (1920, 1080), "#202124" - ) - image_format = "PNG" - draw = ImageDraw.Draw(image) - for row in range(image.height // 20): - draw.text( - (100, row * 20), - f"coordinate (100, {row * 20}) seed={seed} command --verbose " - * 4, - fill="#e8eaed", - ) - draw.rectangle((20, 20, 70, 70), outline="red", width=3) - path = output / f"{name}.{image_format.lower()}" - try: - image.save(path, image_format, **options) - finally: - image.close() - for frame_image in options.get("append_images", []): - frame_image.close() - # Reopen the saved bytes: in-memory Image.format and frame counts - # do not describe what the encoder actually persisted. - with Image.open(path) as saved: - record = { - "path": path.name, - "kind": kind, - "stress": stress, - "seed": seed, - "format": saved.format, - "width": saved.width, - "height": saved.height, - "mode": saved.mode, - "frames": getattr(saved, "n_frames", 1), - "orientation": saved.getexif().get(274, 1), - "has_alpha": saved.mode in ("RGBA", "LA") - or "transparency" in saved.info, - "bytes": path.stat().st_size, - "sha256": hashlib.sha256(path.read_bytes()).hexdigest(), - } - records.append(record) - (output / "manifest.json").write_text( - json.dumps(records, indent=2) + "\n", encoding="utf-8" - ) - return records - - -def main() -> int: - """Generate three paired seeds unless explicitly overridden.""" - parser = argparse.ArgumentParser() - parser.add_argument("output", type=Path) - parser.add_argument("--seeds", nargs="+", type=int, default=[7, 19, 43]) - args = parser.parse_args() - for seed in args.seeds: - records = generate(args.output / str(seed), seed) - print(json.dumps({"seed": seed, "count": len(records)})) - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/scripts/image_memory_bench/history_benchmark.py b/scripts/image_memory_bench/history_benchmark.py deleted file mode 100644 index 0e9563bf6b..0000000000 --- a/scripts/image_memory_bench/history_benchmark.py +++ /dev/null @@ -1,461 +0,0 @@ -"""Measure SQLite history loading and real SDK request bodies for B-only.""" - -# The compact benchmark runner keeps subprocess orchestration visibly linear. -# ruff: noqa: E701, E702 - -from __future__ import annotations - -import argparse -import asyncio -import base64 -import csv -import hashlib -import json -import mimetypes -import os -import subprocess -import sys -import tempfile -import threading -import time -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from pathlib import Path - -import psutil - - -async def prepare_database( - db_path: Path, - media_root: Path, - fixture: Path, - images: int, - distinct: bool, - mode: str, -) -> None: - """Create one SQLite input outside the measured child.""" - sys.path.insert(0, str(Path(__file__).resolve().parents[2])) - from astrbot.core.db.sqlite import SQLiteDatabase - from astrbot.core.utils.image_media_store import ( - ImageMediaStore, - persist_inline_image_refs, - ) - - rows = json.loads(fixture.read_text()) - paths = [fixture.parent / row["path"] for row in rows] - history = [] - raw = [] - for index in range(images): - path = paths[index % len(paths)] if distinct else paths[0] - # Distinct workloads use separately generated, valid image files listed - # by the fixture manifest. Never mutate encoded bytes into invalid images. - data = path.read_bytes() - mime = mimetypes.guess_type(path.name)[0] or "image/png" - raw.append( - { - "type": "image_url", - "image_url": { - "url": f"data:{mime};base64,{base64.b64encode(data).decode()}", - "detail": "high", - }, - } - ) - for index, part in enumerate(raw): - history.extend( - [ - {"role": "user", "content": [part]}, - {"role": "assistant", "content": f"ack {index}"}, - ] - ) - history.extend( - [ - {"role": "user", "content": "final text"}, - {"role": "assistant", "content": "final response"}, - ] - ) - if mode == "reference": - history = persist_inline_image_refs(history, ImageMediaStore(media_root)) - db = SQLiteDatabase(str(db_path)) - await db.initialize() - record = await db.create_conversation( - user_id="bench", platform_id="bench", content=history, title="benchmark" - ) - (db_path.parent / "conversation-id").write_text(record.conversation_id) - await db.engine.dispose() - - -def count_images(value: object) -> int: - """Count data image URLs in a decoded SDK body.""" - if isinstance(value, dict): - return sum(count_images(item) for item in value.values()) - if isinstance(value, list): - return sum(count_images(item) for item in value) - return int(isinstance(value, str) and value.startswith("data:image/")) - - -async def run_child(args: argparse.Namespace) -> dict: - """Read through SQLite and ConversationManager before selecting context.""" - sys.path.insert(0, str(args.repo.resolve())) - from openai import AsyncOpenAI - - from astrbot.core.agent.context.config import ContextConfig - from astrbot.core.agent.context.manager import ContextManager - from astrbot.core.agent.message import Message - from astrbot.core.conversation_mgr import ConversationManager - from astrbot.core.db.sqlite import SQLiteDatabase - from astrbot.core.utils.image_media_store import ( - ImageMediaStore, - materialize_image_media_refs, - ) - - phases = {} - - def mark(phase: str) -> None: - with args.marker.open("a", encoding="utf-8") as marker: - marker.write( - json.dumps({"phase": phase, "monotonic": time.monotonic()}) + "\n" - ) - - mark("db_and_convert") - started = time.perf_counter() - db = SQLiteDatabase(str(args.database)) - await db.initialize() - conversation = await ConversationManager(db).get_conversation( - "bench", args.conversation_id - ) - phases["db_and_convert_ms"] = (time.perf_counter() - started) * 1000 - mark("json_load_bind") - started = time.perf_counter() - history = json.loads(conversation.history) - messages = [Message.model_validate(item) for item in history] - phases["json_load_bind_ms"] = (time.perf_counter() - started) * 1000 - mark("selection") - started = time.perf_counter() - selected = await ContextManager( - ContextConfig(enforce_max_turns=args.keep_turns) - ).process(messages) - phases["selection_ms"] = (time.perf_counter() - started) * 1000 - mark("materialize") - selected = await materialize_image_media_refs( - selected, ImageMediaStore(args.media_root) - ) - payload = [message.model_dump() for message in selected] - mark("sdk") - async with AsyncOpenAI( - api_key="local", base_url=args.endpoint, max_retries=0 - ) as client: - for _ in range(args.warmup): - await client.chat.completions.create(model="local", messages=payload) - await client.chat.completions.create(model="local", messages=payload) - await db.engine.dispose() - return { - "phases_ms": phases, - "selected_messages": len(payload), - "selected_images": count_images(payload), - } - - -def main() -> int: - """Prepare paired databases and monitor fresh children.""" - parser = argparse.ArgumentParser() - parser.add_argument("fixture", type=Path) - parser.add_argument("output", type=Path) - parser.add_argument( - "--repo", type=Path, default=Path(__file__).resolve().parents[2] - ) - parser.add_argument("--images", type=int, default=50) - parser.add_argument("--keep-turns", type=int, default=25) - parser.add_argument("--repeats", type=int, default=10) - parser.add_argument("--cold-repeats", type=int) - parser.add_argument("--warm-repeats", type=int) - parser.add_argument("--warmup", type=int, default=3) - parser.add_argument("--distinct", action="store_true") - parser.add_argument("--csv", type=Path) - parser.add_argument("--rss-jsonl", type=Path) - parser.add_argument("--child", action="store_true") - parser.add_argument("--database", type=Path) - parser.add_argument("--media-root", type=Path) - parser.add_argument("--history-mode", choices=("inline", "reference")) - parser.add_argument("--conversation-id", default="") - parser.add_argument("--endpoint") - parser.add_argument("--body-file", type=Path) - parser.add_argument("--marker", type=Path) - args = parser.parse_args() - cold_repeats = args.cold_repeats if args.cold_repeats is not None else args.repeats - warm_repeats = args.warm_repeats if args.warm_repeats is not None else args.repeats - if min(args.repeats, cold_repeats, warm_repeats, args.images, args.keep_turns) < 1: - parser.error("repetition, image, and context-window values must be positive") - fixture_rows = json.loads(args.fixture.read_text()) - max_fixture_bytes = max( - ( - row.get("bytes") or row.get("prepared_bytes") or row.get("source_bytes", 0) - for row in fixture_rows - ), - default=0, - ) - budget_class = ( - "over-4MiB-source-stress" - if 4 * ((max_fixture_bytes + 2) // 3) > 4 * 1024 * 1024 - else "representative-within-4MiB-source" - ) - if args.child: - if ( - args.database is None - or args.media_root is None - or args.endpoint is None - or args.body_file is None - or args.marker is None - or args.history_mode is None - ): - parser.error( - "child requires database, media root, endpoint, body file, marker, and history mode" - ) - result = asyncio.run(run_child(args)) - if args.body_file.exists(): - metadata = json.loads(args.body_file.read_text()) - result.update( - { - "http_json_bytes": metadata["bytes"], - "http_body_sha256": metadata["sha256"], - "http_images": metadata["images"], - } - ) - else: - result.update( - { - "http_json_bytes": None, - "http_body_sha256": None, - "http_images": 0, - } - ) - if not result["selected_images"] or not result["http_images"]: - raise RuntimeError("actual SDK body contains no images") - args.output.write_text(json.dumps(result) + "\n") - return 0 - args.output.parent.mkdir(parents=True, exist_ok=True) - args.output.write_text("", encoding="utf-8") - csv_file = args.csv.open("w", newline="", encoding="utf-8") if args.csv else None - writer = None - received_body = {"path": None} - rss_path = args.output.with_name(args.output.stem + ".rss.jsonl") - if args.rss_jsonl: - rss_path = args.rss_jsonl - rss_path.parent.mkdir(parents=True, exist_ok=True) - rss_file = rss_path.open("w", encoding="utf-8") - - class Handler(BaseHTTPRequestHandler): - def do_POST(self): - data = self.rfile.read(int(self.headers.get("Content-Length", "0"))) - body = json.loads(data) - Path(received_body["path"]).write_text( - json.dumps( - { - "bytes": len(data), - "sha256": hashlib.sha256(data).hexdigest(), - "images": count_images(body), - } - ) - ) - response = b'{"choices":[{"message":{"content":"ok"}}]}' - self.send_response(200) - self.send_header("Content-Length", str(len(response))) - self.end_headers() - self.wfile.write(response) - - def log_message(self, *_args): - pass - - server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) - threading.Thread(target=server.serve_forever, daemon=True).start() - endpoint = f"http://127.0.0.1:{server.server_port}/v1" - try: - for mode in ("inline", "reference"): - for repeat in range(cold_repeats + warm_repeats): - with tempfile.TemporaryDirectory( - prefix="astrbot-history-input-" - ) as temp: - root = Path(temp) - db_path = root / f"{mode}.db" - media = root / mode - asyncio.run( - prepare_database( - db_path, - media, - args.fixture, - args.images, - args.distinct, - mode, - ) - ) - conversation_id = (db_path.parent / "conversation-id").read_text() - child_output = root / "child.jsonl" - body_file = root / "http-metadata.json" - received_body["path"] = body_file - warmup = 0 if repeat < cold_repeats else args.warmup - command = [ - sys.executable, - str(Path(__file__).resolve()), - str(args.fixture), - str(child_output), - "--child", - "--repo", - str(args.repo), - "--database", - str(db_path), - "--media-root", - str(media), - "--conversation-id", - conversation_id, - "--keep-turns", - str(args.keep_turns), - "--history-mode", - mode, - "--warmup", - str(warmup), - "--endpoint", - endpoint, - "--body-file", - str(body_file), - "--marker", - str(root / "marker.jsonl"), - ] - proc = subprocess.Popen( - command, - cwd=args.repo, - env={**os.environ, "ASTRBOT_ROOT": str(root)}, - ) - monitor = psutil.Process(proc.pid) - samples = [] - started = time.monotonic() - limit = min(2 * 1024**3, psutil.virtual_memory().available // 4) - stop_reason = None - while proc.poll() is None: - try: - now = time.monotonic() - rss = monitor.memory_info().rss - private_bytes = None - try: - private_bytes = monitor.memory_full_info().uss - except (AttributeError, psutil.AccessDenied): - pass - samples.append( - { - "seconds": now - started, - "monotonic": now, - "rss": rss, - "private_bytes": private_bytes, - } - ) - if now - started > 120: - stop_reason = "timeout" - elif rss > limit: - stop_reason = "memory_limit" - if stop_reason: - proc.kill() - break - except psutil.NoSuchProcess: - break - time.sleep(0.01) - proc.wait() - if child_output.exists(): - result = json.loads(child_output.read_text()) - else: - result = { - "status": "error", - "error": stop_reason or "child_failed", - } - marker_file = root / "marker.jsonl" - markers = ( - [ - json.loads(line) - for line in marker_file.read_text().splitlines() - ] - if marker_file.exists() - else [] - ) - for sample in samples: - eligible = [ - marker - for marker in markers - if marker["monotonic"] <= sample["monotonic"] - ] - sample["phase"] = ( - max(eligible, key=lambda marker: marker["monotonic"])[ - "phase" - ] - if eligible - else "pre-start" - ) - result.update( - { - "mode": mode, - "repeat": repeat, - "cold": repeat < cold_repeats, - "warmup_mode": "sdk_warmup" if warmup else "cold", - "warmup_scope": "sdk_send_only", - "workload": f"sqlite-history-{args.images}-{'distinct' if args.distinct else 'same'}", - "budget_class": budget_class, - "rss_high_water_bytes": max( - (sample["rss"] for sample in samples), default=None - ), - "private_high_water_bytes": max( - ( - sample["private_bytes"] - for sample in samples - if sample["private_bytes"] is not None - ), - default=None, - ), - "stop_reason": stop_reason, - "rss_sample_count": len(samples), - "rss_jsonl": str(rss_path), - "returncode": proc.returncode, - "status": ( - "ok" - if proc.returncode == 0 and stop_reason is None - else "error" - ), - } - ) - for sample in samples: - rss_file.write( - json.dumps( - { - "mode": mode, - "repeat": repeat, - **sample, - } - ) - + "\n" - ) - rss_file.flush() - with args.output.open("a", encoding="utf-8") as stream: - stream.write(json.dumps(result) + "\n") - if csv_file: - fields = [ - "mode", - "repeat", - "cold", - "workload", - "rss_high_water_bytes", - "private_high_water_bytes", - "rss_sample_count", - "status", - "returncode", - "http_json_bytes", - "http_body_sha256", - "http_images", - ] - if writer is None: - writer = csv.DictWriter(csv_file, fieldnames=fields) - writer.writeheader() - writer.writerow({field: result.get(field) for field in fields}) - finally: - server.shutdown() - rss_file.close() - if csv_file: - csv_file.close() - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/scripts/image_memory_bench/lifecycle_workloads.py b/scripts/image_memory_bench/lifecycle_workloads.py deleted file mode 100644 index 75c80846e0..0000000000 --- a/scripts/image_memory_bench/lifecycle_workloads.py +++ /dev/null @@ -1,371 +0,0 @@ -"""Run real history lifecycle workloads in an isolated, measured child process.""" - -from __future__ import annotations - -import argparse -import asyncio -import base64 -import hashlib -import json -import subprocess -import sys -import tempfile -import threading -import time -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from pathlib import Path -from typing import Any - -import psutil - -REQUESTS = 200 -SESSIONS = 4 -WINDOW_TURNS = 25 - - -def _count_images(value: Any) -> int: - if isinstance(value, dict): - return sum(_count_images(item) for item in value.values()) - if isinstance(value, list): - return sum(_count_images(item) for item in value) - return int(isinstance(value, str) and value.startswith("data:image/")) - - -def _image_hashes(value: Any) -> list[str]: - if isinstance(value, dict): - if value.get("type") == "image_url" and isinstance( - value.get("image_url"), dict - ): - encoded = value["image_url"].get("url", "").split(",", 1)[-1] - try: - return [hashlib.sha256(base64.b64decode(encoded)).hexdigest()] - except Exception: - return [] - return [item for child in value.values() for item in _image_hashes(child)] - if isinstance(value, list): - return [item for child in value for item in _image_hashes(child)] - return [] - - -async def _child(args: argparse.Namespace) -> dict[str, Any]: - """Execute the production database, context, materialization, and SDK path.""" - sys.path.insert(0, str(args.repo.resolve())) - from openai import AsyncOpenAI - - from astrbot.core.agent.context.config import ContextConfig - from astrbot.core.agent.context.manager import ContextManager - from astrbot.core.agent.message import Message - from astrbot.core.conversation_mgr import ConversationManager - from astrbot.core.db.sqlite import SQLiteDatabase - from astrbot.core.utils.image_media_store import ( - ImageMediaStore, - materialize_image_media_refs, - ) - - async def load_payload( - database: SQLiteDatabase, conversation_id: str - ) -> tuple[list, str, list]: - conversation = await ConversationManager(database).get_conversation( - "bench", conversation_id - ) - messages = [ - Message.model_validate(item) for item in json.loads(conversation.history) - ] - selected = await ContextManager( - ContextConfig(enforce_max_turns=args.window_turns) - ).process(messages) - selected = await materialize_image_media_refs( - selected, ImageMediaStore(args.media_root) - ) - payload = [message.model_dump() for message in selected] - semantic_hash = hashlib.sha256( - json.dumps(payload, sort_keys=True, ensure_ascii=False).encode() - ).hexdigest() - return payload, semantic_hash, messages - - async def session( - client: AsyncOpenAI, session_id: str, request_count: int - ) -> dict[str, Any]: - database = SQLiteDatabase(str(args.database)) - await database.initialize() - semantic_hashes: list[str] = [] - image_counts: list[int] = [] - image_hash_orders: list[list[str]] = [] - for turn in range(request_count): - _, _, messages = await load_payload(database, session_id) - image_refs = [ - part.model_dump() if hasattr(part, "model_dump") else part - for message in messages - for part in ( - message.content if isinstance(message.content, list) else [] - ) - if (part.model_dump() if hasattr(part, "model_dump") else part).get( - "type" - ) - == "image_media_ref" - ] - if not image_refs: - raise RuntimeError( - "lifecycle fixture has no persisted image references" - ) - messages.extend( - [ - Message(role="user", content=[image_refs[turn % len(image_refs)]]), - Message(role="assistant", content=f"ack {turn}"), - ] - ) - await ConversationManager(database).update_conversation( - "bench", - session_id, - history=[message.model_dump() for message in messages], - ) - payload, semantic_hash, _ = await load_payload(database, session_id) - semantic_hashes.append(semantic_hash) - image_counts.append(_count_images(payload)) - image_hash_orders.append(_image_hashes(payload)) - if not image_hash_orders[-1] or image_counts[-1] > args.window_turns: - raise RuntimeError("active image window is empty or unbounded") - await client.chat.completions.create(model="local", messages=payload) - del payload - await database.engine.dispose() - return { - "session": session_id, - "requests": request_count, - "image_counts": image_counts, - "image_hash_orders": image_hash_orders, - "semantic_hashes": semantic_hashes, - } - - if args.restart_check: - checks = [] - database = SQLiteDatabase(str(args.database)) - await database.initialize() - for session_id in args.conversation_ids: - payload, semantic_hash, _ = await load_payload(database, session_id) - checks.append( - { - "session": session_id, - "images": _count_images(payload), - "semantic_hash": semantic_hash, - "image_hash_order": _image_hashes(payload), - } - ) - await database.engine.dispose() - return {"restart_checks": checks} - - async with AsyncOpenAI( - api_key="local", base_url=args.endpoint, max_retries=0 - ) as client: - session_count = len(args.conversation_ids) - base_requests, remainder = divmod(args.request_count, session_count) - request_counts = [ - base_requests + (index < remainder) for index in range(session_count) - ] - results = await asyncio.gather( - *( - session(client, cid, request_count) - for cid, request_count in zip(args.conversation_ids, request_counts) - ) - ) - return { - "sessions": results, - "total_requests": sum(request_counts), - "request_count_per_session": request_counts, - "window_turns": args.window_turns, - } - - -def main() -> int: - parser = argparse.ArgumentParser() - parser.add_argument("output", type=Path) - parser.add_argument( - "--repo", type=Path, default=Path(__file__).resolve().parents[2] - ) - parser.add_argument("--child", action="store_true") - parser.add_argument("--database", type=Path) - parser.add_argument("--media-root", type=Path) - parser.add_argument( - "--conversation-id", dest="conversation_ids", action="append", default=[] - ) - parser.add_argument("--endpoint") - parser.add_argument("--rss-jsonl", type=Path) - parser.add_argument("--restart-check", action="store_true") - parser.add_argument("--session-count", type=int, choices=(1, 4), default=4) - parser.add_argument("--request-count", type=int, default=REQUESTS) - parser.add_argument("--window-turns", type=int, default=WINDOW_TURNS) - args = parser.parse_args() - if args.request_count < 1 or args.window_turns < 1: - parser.error("request count and window turns must be positive") - if args.child: - if ( - args.database is None - or args.media_root is None - or args.endpoint is None - or len(args.conversation_ids) != args.session_count - ): - parser.error( - "child requires database, media root, endpoint, and the requested conversation IDs" - ) - result = asyncio.run(_child(args)) - args.output.write_text(json.dumps(result) + "\n", encoding="utf-8") - return 0 - if len(args.conversation_ids) != args.session_count: - parser.error( - f"exactly {args.session_count} --conversation-id values are required" - ) - if args.database is None or args.media_root is None: - parser.error("--database and --media-root are required") - - requests: list[dict[str, Any]] = [] - - class Sink(BaseHTTPRequestHandler): - def do_POST(self) -> None: - size = int(self.headers.get("Content-Length", "0")) - body = self.rfile.read(size) - requests.append( - { - "bytes": len(body), - "sha256": hashlib.sha256(body).hexdigest(), - "images": _count_images(json.loads(body)), - } - ) - response = b'{"choices":[{"message":{"content":"ok"}}]}' - self.send_response(200) - self.send_header("Content-Length", str(len(response))) - self.end_headers() - self.wfile.write(response) - - def log_message(self, *_args: Any) -> None: - pass - - server = ThreadingHTTPServer(("127.0.0.1", 0), Sink) - threading.Thread(target=server.serve_forever, daemon=True).start() - with tempfile.TemporaryDirectory(prefix="astrbot-lifecycle-") as run_dir: - child_output = Path(run_dir) / "child.json" - rss_path = args.rss_jsonl or args.output.with_name( - args.output.stem + ".rss.jsonl" - ) - command = [ - sys.executable, - str(Path(__file__).resolve()), - str(child_output), - "--child", - "--repo", - str(args.repo.resolve()), - "--database", - str(args.database.resolve()), - "--media-root", - str(args.media_root.resolve()), - "--endpoint", - f"http://127.0.0.1:{server.server_port}/v1", - "--session-count", - str(args.session_count), - "--request-count", - str(args.request_count), - "--window-turns", - str(args.window_turns), - ] - for conversation_id in args.conversation_ids: - command.extend(["--conversation-id", conversation_id]) - started = time.monotonic() - process = subprocess.Popen(command, cwd=args.repo) - child = psutil.Process(process.pid) - stop_reason = None - with rss_path.open("w", encoding="utf-8") as rss: - while process.poll() is None: - current = child.memory_info().rss - private_bytes = None - try: - private_bytes = child.memory_full_info().uss - except (AttributeError, psutil.AccessDenied): - pass - rss.write( - json.dumps( - { - "seconds": time.monotonic() - started, - "rss": current, - "private_bytes": private_bytes, - } - ) - + "\n" - ) - rss.flush() - if ( - current > min(2 * 1024**3, psutil.virtual_memory().available // 4) - or time.monotonic() - started > 120 - ): - stop_reason = ( - "memory_limit" - if current - > min(2 * 1024**3, psutil.virtual_memory().available // 4) - else "timeout" - ) - process.kill() - break - time.sleep(0.01) - return_code = process.wait() - server.shutdown() - result = ( - json.loads(child_output.read_text()) - if child_output.exists() - else {"error": "child failed"} - ) - validation_errors = [] - restart_result = {"skipped": stop_reason or "child_failed"} - if stop_reason is None and return_code == 0: - restart_output = Path(run_dir) / "restart.json" - restart_command = command.copy() - restart_command[2] = str(restart_output) - restart_command.append("--restart-check") - restart_process = subprocess.run( - restart_command, cwd=args.repo, check=False - ) - restart_result = ( - json.loads(restart_output.read_text()) - if restart_output.exists() - else {"error": "restart child failed"} - ) - if restart_result.get("restart_checks"): - expected = { - item["session"]: item for item in result.get("sessions", []) - } - for check in restart_result["restart_checks"]: - prior = expected.get(check["session"]) - if ( - prior is None - or check["images"] != prior["image_counts"][-1] - or check["image_hash_order"] != prior["image_hash_orders"][-1] - ): - validation_errors.append( - "fresh-process restart changed the active image window" - ) - if restart_process.returncode != 0: - return_code = restart_process.returncode - result.update( - { - "returncode": return_code, - "stop_reason": stop_reason, - "request_count": len(requests), - "expected_request_count": args.request_count, - "requests": requests, - "rss_jsonl": str(rss_path), - "restart": restart_result, - } - ) - if any( - request["images"] <= 0 or request["images"] > args.window_turns - for request in requests - ): - validation_errors.append("sink observed an invalid active image window") - if len(requests) != args.request_count: - validation_errors.append("sink did not observe the expected request count") - if validation_errors: - result["validation_errors"] = validation_errors - return_code = return_code or 1 - result["returncode"] = return_code - args.output.write_text(json.dumps(result) + "\n", encoding="utf-8") - return return_code - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/scripts/image_memory_bench/measure.py b/scripts/image_memory_bench/measure.py deleted file mode 100644 index 74c4bdb26e..0000000000 --- a/scripts/image_memory_bench/measure.py +++ /dev/null @@ -1,399 +0,0 @@ -"""Measure real image preparation and SDK requests in isolated processes.""" - -from __future__ import annotations - -import argparse -import asyncio -import base64 -import csv -import hashlib -import json -import os -import subprocess -import sys -import tempfile -import threading -import time -import tracemalloc -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from pathlib import Path - -import psutil - - -def _base64_payload_metrics(encoded: str) -> tuple[int, str]: - """Hash an encoded payload in chunks without retaining decoded image bytes. - - Args: - encoded: Base64 payload without the data-URI prefix. - - Returns: - The decoded byte count and SHA-256 digest. - """ - digest = hashlib.sha256() - decoded_bytes = 0 - remainder = "" - chunk_size = 4 * 1024 * 1024 - for start in range(0, len(encoded), chunk_size): - chunk = remainder + encoded[start : start + chunk_size] - usable = len(chunk) - len(chunk) % 4 - if usable: - decoded = base64.b64decode(chunk[:usable]) - digest.update(decoded) - decoded_bytes += len(decoded) - remainder = chunk[usable:] - if remainder: - decoded = base64.b64decode(remainder) - digest.update(decoded) - decoded_bytes += len(decoded) - return decoded_bytes, digest.hexdigest() - - -async def run_child(args: argparse.Namespace) -> dict: - """Execute production preparation and request assembly in the chosen checkout. - - Args: - args: Explicit source, checkout, local server and diagnostic options. - - Returns: - Sizes, phase durations and resource counters without image content. - """ - sys.path.insert(0, str(args.repo.resolve())) - from openai import AsyncOpenAI - - from astrbot.core.provider.entities import ProviderRequest - - if args.trace_python: - tracemalloc.start() - process = psutil.Process() - baseline_rss = process.memory_info().rss - phases: dict[str, float] = {} - stage = "startup" - phase_marks = [(time.monotonic(), stage)] - failure: dict[str, str] | None = None - source_bytes = args.image.stat().st_size - source_digest = hashlib.sha256() - with args.image.open("rb") as source_stream: - for chunk in iter(lambda: source_stream.read(1024 * 1024), b""): - source_digest.update(chunk) - source_sha256 = source_digest.hexdigest() - prepared_bytes: int | None = None - prepared_sha256: str | None = None - started = time.perf_counter() - with tempfile.TemporaryDirectory(prefix="astrbot-image-runtime-") as runtime: - os.environ["ASTRBOT_ROOT"] = runtime - before_files = set(Path(runtime).rglob("*")) - try: - stage = "provider_prepare_and_assemble" - phase_marks.append((time.monotonic(), stage)) - tick = time.perf_counter() - request = ProviderRequest( - prompt="Describe this image.", - image_urls=[str(args.image.resolve())], - ) - context = await request.assemble_context() - phases["provider_prepare_and_assemble_ms"] = ( - time.perf_counter() - tick - ) * 1000 - image_part = next( - part - for part in context.get("content", []) - if part.get("type") == "image_url" - ) - data_url = image_part["image_url"]["url"] - stage = "benchmark_payload_verification" - phase_marks.append((time.monotonic(), stage)) - prepared_bytes, prepared_sha256 = _base64_payload_metrics( - data_url.split(",", 1)[1] - ) - stage = "sdk_send" - phase_marks.append((time.monotonic(), stage)) - async with AsyncOpenAI( - api_key="local-experiment-only", - base_url=args.endpoint, - max_retries=0, - ) as client: - for _ in range(args.warmup): - await client.chat.completions.create( - model="local", messages=[context] - ) - tick = time.perf_counter() - for _ in range(args.repeat): - await client.chat.completions.create( - model="local", messages=[context] - ) - phases["sdk_send_ms"] = (time.perf_counter() - tick) * 1000 - del context, request - except Exception as exc: - failure = {"stage": stage, "error_type": type(exc).__name__} - finally: - residual = [ - path - for path in Path(runtime).rglob("*") - if path.is_file() and path not in before_files - ] - result = { - "source_bytes": source_bytes, - "prepared_bytes": prepared_bytes, - "source_sha256": source_sha256, - "prepared_sha256": prepared_sha256, - "preserved_source_bytes": ( - prepared_sha256 == source_sha256 - if prepared_sha256 is not None - else None - ), - "base64_bytes": ( - 4 * ((prepared_bytes + 2) // 3) - if prepared_bytes is not None - else None - ), - "phase_ms": phases, - "phase_marks": phase_marks, - "last_stage": stage, - "failure": failure, - "baseline_rss": baseline_rss, - "post_request_rss": process.memory_info().rss, - "temporary_files_remaining": len(residual), - "temporary_bytes_remaining": sum( - path.stat().st_size for path in residual - ), - "elapsed_ms": (time.perf_counter() - started) * 1000, - } - if args.trace_python: - result["python_current_bytes"], result["python_peak_bytes"] = ( - tracemalloc.get_traced_memory() - ) - tracemalloc.stop() - if sys.platform != "win32": - import resource - - peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss - result["os_peak_rss_bytes"] = ( - peak if sys.platform == "darwin" else peak * 1024 - ) - return result - - -def main() -> int: - """Run an isolated request and record measurements or a bounded failure.""" - parser = argparse.ArgumentParser() - parser.add_argument("image", type=Path) - parser.add_argument("output", type=Path) - parser.add_argument( - "--repo", type=Path, default=Path(__file__).resolve().parents[2] - ) - parser.add_argument("--trace-python", action="store_true") - parser.add_argument("--child", action="store_true") - parser.add_argument("--endpoint") - parser.add_argument("--repeat", type=int, default=1) - parser.add_argument("--warmup", type=int, default=0) - parser.add_argument("--csv") - parser.add_argument("--workload", default="single-image") - args = parser.parse_args() - if args.child: - try: - result = asyncio.run(run_child(args)) - result["status"] = "error" if result.get("failure") else "ok" - except Exception as error: - result = {"status": "error", "error_type": type(error).__name__} - args.output.write_text(json.dumps(result), encoding="utf-8") - return 0 if result["status"] == "ok" else 1 - - # The server measures actual SDK bytes without storing the request body. - request_sizes = [] - - class Handler(BaseHTTPRequestHandler): - def do_POST(self): - length = int(self.headers.get("Content-Length", "0")) - remaining = length - while remaining: - chunk = self.rfile.read(min(65536, remaining)) - if not chunk: - break - remaining -= len(chunk) - request_sizes.append(length - remaining) - body = b'{"id":"local","object":"chat.completion","created":0,"model":"local","choices":[{"index":0,"message":{"role":"assistant","content":"fixture response"},"finish_reason":"stop"}],"usage":{"prompt_tokens":100,"completion_tokens":2,"total_tokens":102}}' - self.send_response(200) - self.send_header("Content-Type", "application/json") - self.send_header("Content-Length", str(len(body))) - self.end_headers() - self.wfile.write(body) - - def log_message(self, *args): - pass - - server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) - server_thread = threading.Thread(target=server.serve_forever, daemon=True) - server_thread.start() - limit = min(2 * 1024**3, psutil.virtual_memory().available // 4) - started = time.monotonic() - samples = [] - stop_reason = None - try: - with tempfile.TemporaryDirectory(prefix="astrbot-image-measure-") as run_dir: - child_result = Path(run_dir) / "child.json" - command = [ - sys.executable, - str(Path(__file__).resolve()), - str(args.image.resolve()), - str(child_result), - "--child", - "--repo", - str(args.repo.resolve()), - "--endpoint", - f"http://127.0.0.1:{server.server_port}/v1", - ] - if args.repeat != 1: - command.extend(["--repeat", str(args.repeat)]) - if args.warmup: - command.extend(["--warmup", str(args.warmup)]) - command.extend(["--workload", args.workload]) - if args.trace_python: - command.append("--trace-python") - env = os.environ.copy() - env["ASTRBOT_ROOT"] = str(Path(run_dir) / "runtime") - with ( - (Path(run_dir) / "stdout").open("wb") as out, - (Path(run_dir) / "stderr").open("wb") as err, - ): - proc = subprocess.Popen( - command, cwd=args.repo, env=env, stdout=out, stderr=err - ) - process = psutil.Process(proc.pid) - try: - while proc.poll() is None: - try: - children = process.children(recursive=True) - rss = process.memory_info().rss - descendants_rss = sum( - p.memory_info().rss for p in children if p.is_running() - ) - info = process.memory_info() - samples.append( - { - "seconds": time.monotonic() - started, - "rss": rss, - "descendants_rss": descendants_rss, - "private_bytes": getattr(info, "private", None), - } - ) - if rss + descendants_rss > limit: - stop_reason = "memory_limit" - elif time.monotonic() - started > 120: - stop_reason = "timeout" - if stop_reason: - for child in children: - child.kill() - proc.kill() - break - except psutil.NoSuchProcess: - break - time.sleep(0.01) - finally: - if proc.poll() is None: - proc.kill() - proc.wait() - result = ( - json.loads(child_result.read_text()) - if child_result.exists() - else {"status": "error"} - ) - phase_peaks = {} - for sample in samples: - phase = "process_startup" - for timestamp, name in result.get("phase_marks", []): - if timestamp > started + sample["seconds"]: - break - phase = name - sample["phase"] = phase - phase_peaks[phase] = max(phase_peaks.get(phase, 0), sample["rss"]) - result.update( - { - "status": stop_reason or result["status"], - "workload": args.workload, - "phase_rss_high_water_bytes": phase_peaks, - "returncode": proc.returncode, - "rss_high_water_bytes": max( - (s["rss"] for s in samples), default=None - ), - "private_high_water_bytes": max( - ( - s["private_bytes"] - for s in samples - if s["private_bytes"] is not None - ), - default=None, - ), - "descendants_high_water_bytes": max( - (s["descendants_rss"] for s in samples), default=None - ), - "request_json_bytes": request_sizes, - "samples": samples, - "trace_python": args.trace_python, - "code_commit": subprocess.check_output( - ["git", "rev-parse", "HEAD"], cwd=args.repo, text=True - ).strip(), - "working_tree_diff_sha256": hashlib.sha256( - subprocess.check_output( - ["git", "diff", "--binary"], cwd=args.repo - ) - ).hexdigest(), - "benchmark_source_sha256": hashlib.sha256( - Path(__file__).read_bytes() - ).hexdigest(), - "monitor_sample_count": len(samples), - } - ) - with args.output.open("a", encoding="utf-8") as output: - output.write(json.dumps(result) + "\n") - if args.csv: - fields = [ - "status", - "workload", - "repeat", - "warmup", - "rss_high_water_bytes", - "private_high_water_bytes", - "post_request_rss", - "temporary_files_remaining", - "temporary_bytes_remaining", - "request_json_bytes", - ] - csv_path = Path(args.csv) - write_header = not csv_path.exists() - with csv_path.open("a", newline="", encoding="utf-8") as csv_file: - writer = csv.DictWriter(csv_file, fieldnames=fields) - if write_header: - writer.writeheader() - writer.writerow( - { - "status": result.get("status"), - "workload": args.workload, - "repeat": args.repeat, - "warmup": args.warmup, - "rss_high_water_bytes": result.get("rss_high_water_bytes"), - "private_high_water_bytes": result.get( - "private_high_water_bytes" - ), - "post_request_rss": result.get("post_request_rss"), - "temporary_files_remaining": result.get( - "temporary_files_remaining" - ), - "temporary_bytes_remaining": result.get( - "temporary_bytes_remaining" - ), - "request_json_bytes": json.dumps( - result.get("request_json_bytes", []) - ), - } - ) - print(json.dumps({k: v for k, v in result.items() if k != "samples"})) - return 0 if result["status"] == "ok" else 1 - finally: - server.shutdown() - server.server_close() - server_thread.join() - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/scripts/image_memory_bench/migrate_history.py b/scripts/image_memory_bench/migrate_history.py deleted file mode 100644 index d2189d7c95..0000000000 --- a/scripts/image_memory_bench/migrate_history.py +++ /dev/null @@ -1,99 +0,0 @@ -"""Explicitly externalize or re-inline exported JSONL conversation histories. - -The input is never modified. Inspect the dry run before using --apply. Keep the -input alongside the media directory as the rollback backup. -""" - -from __future__ import annotations - -import argparse -import asyncio -import json -import sys -from pathlib import Path - -sys.path.insert(0, str(Path(__file__).resolve().parents[2])) - -from astrbot.core.utils.image_media_store import ( - ImageMediaStore, - materialize_image_media_refs, - persist_inline_image_refs, -) - - -async def main() -> int: - """Transform an export into a new file without modifying live histories.""" - parser = argparse.ArgumentParser() - parser.add_argument("input", type=Path) - parser.add_argument("output", type=Path) - parser.add_argument("--media-dir", type=Path, required=True) - parser.add_argument( - "--mode", choices=["externalize", "inline"], default="externalize" - ) - parser.add_argument("--apply", action="store_true") - parser.add_argument("--rollback", action="store_true") - args = parser.parse_args() - if args.input.resolve() == args.output.resolve() or args.output.exists(): - parser.error("Output must be a new file distinct from the rollback input") - if args.rollback: - if not args.apply: - parser.error("--rollback requires --apply") - # The input is the verified original export; rollback writes that exact - # export to a new path, so the original remains the recovery anchor. - args.mode = "inline" - store = ImageMediaStore(args.media_dir) - record_count = 0 - image_count = 0 - output = args.output.open("x", encoding="utf-8") if args.apply else None - try: - with args.input.open(encoding="utf-8") as source: - for line in source: - record = json.loads(line) - raw_history = record.get("history", []) - history = ( - json.loads(raw_history) - if isinstance(raw_history, str) - else raw_history - ) - if not isinstance(history, list): - raise ValueError("History must be a message list") - for message in history: - parts = message.get("content") - if isinstance(parts, list): - image_count += sum( - isinstance(part, dict) - and part.get("type") in {"image_url", "image_media_ref"} - for part in parts - ) - if args.apply: - if args.mode == "externalize" and not args.rollback: - history = persist_inline_image_refs(history, store) - else: - history = await materialize_image_media_refs( - history, store, strict=True - ) - record["history"] = ( - json.dumps(history, ensure_ascii=False) - if isinstance(raw_history, str) - else history - ) - output.write(json.dumps(record, ensure_ascii=False) + "\n") - record_count += 1 - except BaseException: - if output is not None: - output.close() - args.output.unlink(missing_ok=True) - raise - finally: - if output is not None: - output.close() - print( - json.dumps( - {"records": record_count, "images": image_count, "applied": args.apply} - ) - ) - return 0 - - -if __name__ == "__main__": - raise SystemExit(asyncio.run(main())) diff --git a/scripts/image_memory_bench/prepare_lifecycle_fixture.py b/scripts/image_memory_bench/prepare_lifecycle_fixture.py deleted file mode 100644 index b67f9452e5..0000000000 --- a/scripts/image_memory_bench/prepare_lifecycle_fixture.py +++ /dev/null @@ -1,131 +0,0 @@ -"""Prepare a durable multi-session history fixture before measurement.""" - -from __future__ import annotations - -import argparse -import asyncio -import json -from copy import deepcopy -from pathlib import Path - - -async def prepare_fixture( - manifest_path: Path, - output: Path, - session_count: int, - history_turns: int, - include_stress: bool, -) -> dict: - """Create shared media and persisted conversations outside the measured child. - - Args: - manifest_path: Fixture manifest containing valid prepared image paths. - output: Fresh directory receiving the SQLite database and media objects. - session_count: Number of conversations to create. - history_turns: Number of image-bearing turns per conversation. - include_stress: Whether to include high-entropy stress fixtures. - - Returns: - A JSON-compatible manifest consumed by ``lifecycle_workloads.py``. - - Raises: - FileExistsError: The output already contains a database. - ValueError: The input manifest or requested sizes are invalid. - """ - if session_count < 1 or history_turns < 1: - raise ValueError("session count and history turns must be positive") - rows = json.loads(manifest_path.read_text(encoding="utf-8")) - if not isinstance(rows, list) or not rows: - raise ValueError("fixture manifest must contain at least one image") - if not include_stress: - rows = [row for row in rows if not row.get("stress")] - if not rows: - raise ValueError("fixture manifest must contain at least one image") - output.mkdir(parents=True, exist_ok=True) - database_path = output / "history.db" - if database_path.exists(): - raise FileExistsError(f"refusing to overwrite {database_path}") - - import sys - - sys.path.insert(0, str(Path(__file__).resolve().parents[2])) - from astrbot.core.db.sqlite import SQLiteDatabase - from astrbot.core.utils.image_media_store import ImageMediaStore - - media_root = output / "media" - store = ImageMediaStore(media_root) - refs = [] - for row in rows: - path = manifest_path.parent / row["path"] - refs.append(store.put(path.read_bytes(), detail="high")) - - history = [] - for turn in range(history_turns): - history.extend( - [ - { - "role": "user", - "content": [ - {"type": "text", "text": f"fixture turn {turn}"}, - refs[turn % len(refs)].model_dump(), - ], - }, - {"role": "assistant", "content": f"ack {turn}"}, - ] - ) - - database = SQLiteDatabase(str(database_path)) - await database.initialize() - conversation_ids = [] - for index in range(session_count): - conversation = await database.create_conversation( - user_id="bench", - platform_id="bench", - content=deepcopy(history), - title=f"lifecycle fixture {index}", - ) - conversation_ids.append(conversation.conversation_id) - await database.engine.dispose() - return { - "database": str(database_path.resolve()), - "media_root": str(media_root.resolve()), - "conversation_ids": conversation_ids, - "session_count": session_count, - "history_turns": history_turns, - "media_count": len(refs), - } - - -async def main_async(args: argparse.Namespace) -> int: - """Prepare the fixture and write its machine-readable manifest.""" - result = await prepare_fixture( - args.manifest, - args.output, - args.session_count, - args.history_turns, - args.include_stress, - ) - args.output.mkdir(parents=True, exist_ok=True) - output_manifest = args.output / "lifecycle-manifest.json" - output_manifest.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8") - print(json.dumps(result)) - return 0 - - -def main() -> int: - """Parse fixture preparation arguments.""" - parser = argparse.ArgumentParser() - parser.add_argument("manifest", type=Path) - parser.add_argument("output", type=Path) - parser.add_argument("--session-count", type=int, default=4) - parser.add_argument("--history-turns", type=int, default=50) - parser.add_argument( - "--include-stress", - action="store_true", - help="Include high-entropy stress fixtures in the long-running workload.", - ) - return asyncio.run(main_async(parser.parse_args())) - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/scripts/manage_image_history.py b/scripts/manage_image_history.py deleted file mode 100644 index 4a3635bba0..0000000000 --- a/scripts/manage_image_history.py +++ /dev/null @@ -1,434 +0,0 @@ -#!/usr/bin/env python3 -"""Explicit offline migration, repair, rollback, and cleanup for image history.""" - -from __future__ import annotations - -import argparse -import base64 -import hashlib -import json -import shutil -import sqlite3 -import sys -import tempfile -from collections.abc import Iterator -from pathlib import Path -from typing import Any - -from astrbot.core.utils.image_media_store import ImageMediaRef, ImageMediaStore - - -def _digest(path: Path) -> str: - digest = hashlib.sha256() - with path.open("rb") as stream: - for chunk in iter(lambda: stream.read(1024 * 1024), b""): - digest.update(chunk) - return digest.hexdigest() - - -def _logical_digest(connection: sqlite3.Connection) -> str: - """Hash every table schema and row without loading the database in memory.""" - digest = hashlib.sha256() - tables = connection.execute( - "SELECT name, sql, type FROM sqlite_master WHERE type IN ('table','index','trigger','view') ORDER BY name" - ) - for name, sql, object_type in tables: - digest.update(json.dumps([name, sql], ensure_ascii=False).encode()) - if object_type != "table": - continue - rows = connection.execute( - f'SELECT * FROM "{name.replace(chr(34), chr(34) * 2)}"' - ) - for row in rows: - digest.update( - json.dumps( - list(row), ensure_ascii=False, default=str, separators=(",", ":") - ).encode() - ) - return digest.hexdigest() - - -def _manifest_media_name(name: str) -> tuple[str, str]: - path = Path(name) - if path.name != name or path.suffix not in {".bin", ".json"}: - raise ValueError("invalid media manifest filename") - media_id = path.stem - if len(media_id) != 64 or any(char not in "0123456789abcdef" for char in media_id): - raise ValueError("invalid media manifest media id") - return media_id, path.suffix - - -def _inline_parts(value: Any) -> Iterator[tuple[dict[str, Any], str, str]]: - if isinstance(value, list): - for item in value: - yield from _inline_parts(item) - elif isinstance(value, dict): - if value.get("type") == "image_url" and isinstance( - value.get("image_url"), dict - ): - url = value["image_url"].get("url") - if isinstance(url, str) and url.startswith("data:image/"): - header, encoded = url.split(",", 1) - yield value, header[5:].split(";", 1)[0], encoded - else: - for child in value.values(): - yield from _inline_parts(child) - - -def _replace_inline(value: Any, store: ImageMediaStore) -> tuple[Any, int, int]: - changed = 0 - bytes_migrated = 0 - - def visit(item: Any) -> Any: - nonlocal changed, bytes_migrated - if isinstance(item, list): - return [visit(child) for child in item] - if not isinstance(item, dict): - return item - if item.get("type") == "image_url" and isinstance(item.get("image_url"), dict): - image = item["image_url"] - url = image.get("url") - if isinstance(url, str) and url.startswith("data:image/"): - header, encoded = url.split(",", 1) - data = base64.b64decode(encoded, validate=True) - ref = store.put( - data, - header[5:].split(";", 1)[0], - image.get("detail"), - image.get("id"), - ) - changed += 1 - bytes_migrated += len(data) - return ref.model_dump() - return {key: visit(child) for key, child in item.items()} - - return visit(value), changed, bytes_migrated - - -def _connect(path: Path) -> sqlite3.Connection: - connection = sqlite3.connect(path) - connection.row_factory = sqlite3.Row - connection.execute("PRAGMA foreign_keys=ON") - return connection - - -def _backup(db: Path, media: Path, destination: Path) -> Path: - """Create a consistent database snapshot and a verified media snapshot.""" - destination.mkdir(parents=True, exist_ok=False) - snapshot = sqlite3.connect(destination / db.name) - readonly = sqlite3.connect(f"file:{db}?mode=ro", uri=True) - try: - readonly.backup(snapshot) - snapshot.commit() - finally: - readonly.close() - snapshot.close() - if media.exists(): - (destination / "media").mkdir() - manifest_media = [] - for path in media.iterdir(): - if path.is_symlink() or not path.is_file(): - raise RuntimeError("media backup refuses symlinks and non-files") - target = destination / "media" / path.name - shutil.copy2(path, target) - manifest_media.append({"name": path.name, "sha256": _digest(target)}) - else: - manifest_media = [] - manifest = { - "database": _digest(destination / db.name), - "media": manifest_media, - } - if (destination / "media").exists(): - manifest["media"] = sorted(manifest_media, key=lambda item: item["name"]) - (destination / "manifest.json").write_text( - json.dumps(manifest, sort_keys=True) + "\n" - ) - return destination - - -def _migrate(args: argparse.Namespace) -> dict[str, int]: - if args.apply and not args.offline: - raise SystemExit("--apply requires --offline: stop all history writes first") - db = Path(args.database).resolve() - media = Path(args.media).resolve() - connection = _connect(db) - lock = args.apply - try: - if lock: - connection.execute("BEGIN IMMEDIATE") - backup = _backup( - db, - media, - Path(args.backup) - if args.backup - else Path( - tempfile.mkdtemp(prefix="astrbot-image-history-backup-parent-") - ) - / "snapshot", - ) - query = ( - "SELECT inner_conversation_id, conversation_id, content FROM conversations" - ) - parameters: tuple[Any, ...] = () - if args.conversation_id: - query += ( - " WHERE conversation_id IN (" - + ",".join("?" for _ in args.conversation_id) - + ")" - ) - parameters = tuple(args.conversation_id) - total = changed = bytes_migrated = 0 - store = ImageMediaStore(media) - for row in connection.execute(query, parameters): - total += 1 - try: - content = ( - json.loads(row["content"]) - if isinstance(row["content"], str) - else row["content"] - ) - converted, count, amount = ( - _replace_inline(content, store) - if args.apply - else (content, sum(1 for _ in _inline_parts(content)), 0) - ) - except Exception as exc: # noqa: BLE001 - raise RuntimeError( - f"conversation {row['conversation_id']} cannot be parsed or verified" - ) from exc - changed += count - bytes_migrated += amount - if args.apply and count: - connection.execute( - "UPDATE conversations SET content=? WHERE inner_conversation_id=?", - ( - json.dumps(converted, ensure_ascii=False), - row["inner_conversation_id"], - ), - ) - if args.apply: - post_migration_digest = _logical_digest(connection) - connection.commit() - manifest_path = backup / "manifest.json" - manifest = json.loads(manifest_path.read_text()) - manifest["post_migration_logical_digest"] = post_migration_digest - manifest_path.write_text(json.dumps(manifest, sort_keys=True) + "\n") - print( - json.dumps( - { - "backup": str(backup), - "records": total, - "images": changed, - "bytes": bytes_migrated, - } - ) - ) - else: - connection.rollback() - print(json.dumps({"dry_run": True, "records": total, "images": changed})) - return {"records": total, "images": changed} - finally: - connection.close() - - -def _cleanup(args: argparse.Namespace) -> None: - if args.apply and not args.offline: - raise SystemExit("cleanup --apply requires --offline") - connection = _connect(Path(args.database).resolve()) - media = Path(args.media).resolve() - try: - connection.execute("BEGIN IMMEDIATE") - referenced: set[str] = set() - for row in connection.execute("SELECT content FROM conversations"): - try: - content = json.loads(row[0]) if isinstance(row[0], str) else row[0] - if not isinstance(content, list): - raise ValueError("history content is not a list") - for item in content: - if not isinstance(item, dict): - raise ValueError("history message is not an object") - pending = list(content) - while pending: - part = pending.pop() - if isinstance(part, list): - pending.extend(part) - elif isinstance(part, dict): - if part.get("type") == "image_media_ref": - parsed = ImageMediaRef( - media_id=part["media_id"], - mime_type=part["mime_type"], - width=part.get("width"), - height=part.get("height"), - byte_size=part["byte_size"], - detail=part.get("detail"), - version=part.get("version", 1), - image_id=part.get("image_id"), - ) - ImageMediaStore(media).read(parsed, {parsed.media_id}) - referenced.add(parsed.media_id) - else: - pending.extend(part.values()) - except Exception as exc: # noqa: BLE001 - raise RuntimeError("history parse failed; cleanup aborted") from exc - quarantine = Path( - args.quarantine or (media.parent / (media.name + ".quarantine")) - ) - if args.apply: - quarantine.mkdir(parents=True, exist_ok=True) - moved = 0 - for media_id in sorted( - {path.stem for path in media.iterdir() if path.suffix in {".bin", ".json"}} - ): - paths = [media / f"{media_id}.bin", media / f"{media_id}.json"] - if any(path.is_symlink() for path in paths if path.exists()): - continue - if media_id in referenced: - continue - existing_paths = [path for path in paths if path.exists()] - if len(existing_paths) != 2: - continue - try: - metadata = json.loads(paths[1].read_text()) - if ( - metadata.get("media_id") != media_id - or _digest(paths[0]) != media_id - ): - continue - except (OSError, json.JSONDecodeError): - continue - if args.apply: - targets = [quarantine / path.name for path in existing_paths] - if any(target.exists() for target in targets): - raise RuntimeError( - "quarantine target already exists; refusing overwrite" - ) - for path, target in zip(existing_paths, targets): - shutil.move(str(path), target) - moved += len(existing_paths) - if args.apply: - connection.commit() - else: - connection.rollback() - print( - json.dumps( - { - "dry_run": not args.apply, - "quarantined": moved, - "referenced": len(referenced), - } - ) - ) - finally: - connection.close() - - -def _restore(args: argparse.Namespace) -> None: - """Restore only when the database and media are unchanged since migration.""" - if not args.offline: - raise SystemExit("restore requires --offline") - backup = Path(args.restore).resolve() - db = Path(args.database).resolve() - backup_db = backup / db.name - manifest_path = backup / "manifest.json" - if not backup_db.is_file() or not manifest_path.is_file(): - raise SystemExit("backup is incomplete") - manifest = json.loads(manifest_path.read_text()) - if _digest(backup_db) != manifest.get("database"): - raise SystemExit("restore refused: archived database hash mismatch") - archived_media = backup / "media" - manifest_names: set[str] = set() - for item in manifest.get("media", []): - _manifest_media_name(item["name"]) - checksum = item.get("sha256", "") - if ( - item["name"] in manifest_names - or len(checksum) != 64 - or any(char not in "0123456789abcdef" for char in checksum) - ): - raise SystemExit("restore refused: invalid media manifest") - manifest_names.add(item["name"]) - source = archived_media / item["name"] - if ( - source.is_symlink() - or not source.is_file() - or _digest(source) != item["sha256"] - ): - raise SystemExit("restore refused: archived media hash mismatch") - current = _connect(db) - try: - if _logical_digest(current) != manifest.get("post_migration_logical_digest"): - raise SystemExit("restore refused: database changed since migration") - finally: - current.close() - media = Path(args.media).resolve() - for item in manifest.get("media", []): - target = media / item["name"] - if target.exists() and ( - target.is_symlink() or _digest(target) != item["sha256"] - ): - raise SystemExit(f"restore refused: conflicting media: {target.name}") - rollback = db.with_name(db.name + ".before-restore") - if rollback.exists(): - raise SystemExit("refusing to overwrite an existing pre-restore copy") - try: - staged_media: list[tuple[Path, Path]] = [] - staging = Path( - tempfile.mkdtemp(prefix="astrbot-restore-media-", dir=media.parent) - ) - for item in manifest.get("media", []): - source = archived_media / item["name"] - staged = staging / item["name"] - shutil.copy2(source, staged) - if _digest(staged) != item["sha256"]: - raise SystemExit("restore refused: staged media hash mismatch") - staged_media.append((staged, media / item["name"])) - media.mkdir(parents=True, exist_ok=True) - for staged, target in staged_media: - if not target.exists(): - staged.replace(target) - - rollback_connection = sqlite3.connect(rollback) - current = _connect(db) - try: - current.backup(rollback_connection) - rollback_connection.commit() - finally: - current.close() - rollback_connection.close() - destination = _connect(db) - archived = _connect(backup_db) - try: - archived.backup(destination) - destination.commit() - finally: - archived.close() - destination.close() - finally: - if "staging" in locals(): - shutil.rmtree(staging, ignore_errors=True) - print(json.dumps({"restored": str(db), "pre_restore_copy": str(rollback)})) - - -def main() -> int: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("database") - parser.add_argument("--media", required=True) - parser.add_argument("--conversation-id", action="append") - parser.add_argument("--apply", action="store_true") - parser.add_argument("--offline", action="store_true") - parser.add_argument("--backup") - parser.add_argument("--cleanup", action="store_true") - parser.add_argument("--quarantine") - parser.add_argument("--restore") - args = parser.parse_args() - if args.restore: - _restore(args) - elif args.cleanup: - _cleanup(args) - else: - _migrate(args) - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/tests/unit/test_image_lifecycle_workloads.py b/tests/unit/test_image_lifecycle_workloads.py deleted file mode 100644 index 42427d35ef..0000000000 --- a/tests/unit/test_image_lifecycle_workloads.py +++ /dev/null @@ -1,36 +0,0 @@ -from __future__ import annotations - -import subprocess -import sys -from pathlib import Path - - -def test_lifecycle_workload_defines_bounded_matrix_and_no_payload_dump(): - source = Path("scripts/image_memory_bench/lifecycle_workloads.py").read_text( - encoding="utf-8" - ) - assert "REQUESTS = 200" in source - assert "SESSIONS = 4" in source - assert "--request-count" in source - assert "--window-turns" in source - assert "time.sleep(0.01)" in source - assert "120" in source - assert "b64encode" not in source - assert "print(" not in source - - -def test_lifecycle_child_requires_real_database_and_reports_wire_metadata(tmp_path): - output = tmp_path / "result.jsonl" - result = subprocess.run( - [ - sys.executable, - "scripts/image_memory_bench/lifecycle_workloads.py", - str(output), - "--child", - ], - check=False, - capture_output=True, - text=True, - ) - assert result.returncode != 0 - assert not output.exists() diff --git a/tests/unit/test_image_memory_benchmark.py b/tests/unit/test_image_memory_benchmark.py deleted file mode 100644 index 65ec479c97..0000000000 --- a/tests/unit/test_image_memory_benchmark.py +++ /dev/null @@ -1,59 +0,0 @@ -"""Safety checks for the standalone image benchmark utilities.""" - -import json -import subprocess -import sys - - -def test_baseline_manifest_records_fixture_fingerprints(tmp_path): - fixture = tmp_path / "fixtures" / "7" - fixture.mkdir(parents=True) - (fixture / "sample.png").write_bytes(b"fixture") - output = tmp_path / "manifest.json" - subprocess.run( - [ - sys.executable, - "scripts/image_memory_bench/baseline_manifest.py", - str(output), - "--fixtures", - str(tmp_path / "fixtures"), - ], - check=True, - ) - manifest = json.loads(output.read_text()) - assert manifest["fixtures"][0]["sha256"] - - -def test_migration_dry_run_and_rollback_keep_input(tmp_path): - source = tmp_path / "history.jsonl" - source.write_text( - json.dumps({"history": [{"role": "user", "content": "hello"}]}) + "\n" - ) - output = tmp_path / "migrated.jsonl" - media = tmp_path / "media" - command = [ - sys.executable, - "scripts/image_memory_bench/migrate_history.py", - str(source), - str(output), - "--media-dir", - str(media), - ] - subprocess.run(command, check=True) - assert not output.exists() - subprocess.run(command + ["--apply"], check=True) - rollback = tmp_path / "rollback.jsonl" - subprocess.run( - [ - sys.executable, - "scripts/image_memory_bench/migrate_history.py", - str(output), - str(rollback), - "--media-dir", - str(media), - "--rollback", - "--apply", - ], - check=True, - ) - assert output.read_bytes() == rollback.read_bytes() diff --git a/tests/unit/test_manage_image_history.py b/tests/unit/test_manage_image_history.py deleted file mode 100644 index 959db2f520..0000000000 --- a/tests/unit/test_manage_image_history.py +++ /dev/null @@ -1,354 +0,0 @@ -from __future__ import annotations - -import base64 -import importlib.util -import json -import sqlite3 -import subprocess -import sys - -import pytest -from PIL import Image - - -def _image() -> bytes: - import io - - output = io.BytesIO() - Image.new("RGB", (3, 2), "red").save(output, "PNG") - return output.getvalue() - - -def _db(path, content): - connection = sqlite3.connect(path) - connection.execute( - "CREATE TABLE conversations (inner_conversation_id INTEGER PRIMARY KEY, conversation_id TEXT, content JSON)" - ) - connection.execute( - "INSERT INTO conversations VALUES (1, 'one', ?)", (json.dumps(content),) - ) - connection.commit() - connection.close() - - -def _run(db, media, *args): - return subprocess.run( - [ - sys.executable, - "scripts/manage_image_history.py", - str(db), - "--media", - str(media), - *args, - ], - check=False, - capture_output=True, - text=True, - ) - - -def test_migration_defaults_to_dry_run(tmp_path): - data = _image() - history = [ - { - "role": "user", - "content": [ - { - "type": "image_url", - "image_url": { - "url": "data:image/png;base64," - + base64.b64encode(data).decode() - }, - } - ], - } - ] - db = tmp_path / "db.sqlite" - media = tmp_path / "media" - _db(db, history) - result = _run(db, media) - assert result.returncode == 0 - assert not media.exists() - assert ( - json.loads( - sqlite3.connect(db) - .execute("SELECT content FROM conversations") - .fetchone()[0] - ) - == history - ) - - -def test_apply_creates_backup_and_references(tmp_path): - data = _image() - history = [ - { - "role": "user", - "content": [ - { - "type": "image_url", - "image_url": { - "url": "data:image/png;base64," - + base64.b64encode(data).decode(), - "id": "keep", - }, - } - ], - } - ] - db = tmp_path / "db.sqlite" - media = tmp_path / "media" - backup = tmp_path / "backup" - _db(db, history) - result = _run(db, media, "--apply", "--offline", "--backup", str(backup)) - assert result.returncode == 0, result.stderr - content = json.loads( - sqlite3.connect(db).execute("SELECT content FROM conversations").fetchone()[0] - ) - ref = content[0]["content"][0] - assert ref["type"] == "image_media_ref" and ref["image_id"] == "keep" - assert (backup / "db.sqlite").exists() and (backup / "manifest.json").exists() - - -def test_apply_requires_offline_and_cleanup_dry_run_is_read_only(tmp_path): - db = tmp_path / "db.sqlite" - media = tmp_path / "media" - _db(db, []) - media.mkdir() - (media / "dead.bin").write_bytes(b"x") - result = _run(db, media, "--apply") - assert result.returncode != 0 - result = _run(db, media, "--cleanup") - assert result.returncode == 0 - assert (media / "dead.bin").exists() - - -def test_cleanup_aborts_on_malformed_history(tmp_path): - db = tmp_path / "db.sqlite" - media = tmp_path / "media" - _db(db, {"not": "a history"}) - media.mkdir() - (media / ("a" * 64 + ".bin")).write_bytes(b"orphan") - result = _run(db, media, "--cleanup", "--apply", "--offline") - assert result.returncode != 0 - assert (media / ("a" * 64 + ".bin")).exists() - - -def test_cleanup_preserves_references_nested_in_retained_messages(tmp_path): - from astrbot.core.utils.image_media_store import ImageMediaStore - - media = tmp_path / "media" - store = ImageMediaStore(media) - ref = store.put(_image()) - db = tmp_path / "db.sqlite" - _db(db, [{"role": "tool", "content": [{"resource": ref.model_dump()}]}]) - result = _run(db, media, "--cleanup", "--apply", "--offline") - assert result.returncode == 0, result.stderr - assert store.read(ref, {ref.media_id}) == _image() - - -def test_cleanup_keeps_shared_media_until_last_history_reference_is_removed(tmp_path): - from astrbot.core.utils.image_media_store import ImageMediaStore - - media = tmp_path / "media" - store = ImageMediaStore(media) - ref = store.put(_image()) - first_history = [{"role": "user", "content": [ref.model_dump()]}] - second_history = [{"role": "user", "content": [ref.model_dump()]}] - db = tmp_path / "db.sqlite" - _db(db, first_history) - connection = sqlite3.connect(db) - connection.execute( - "INSERT INTO conversations VALUES (2, 'two', ?)", - (json.dumps(second_history),), - ) - connection.commit() - connection.close() - - first_delete = _run(db, media, "--cleanup", "--apply", "--offline") - assert first_delete.returncode == 0, first_delete.stderr - assert (media / f"{ref.media_id}.bin").exists() - - connection = sqlite3.connect(db) - connection.execute("DELETE FROM conversations WHERE inner_conversation_id=1") - connection.commit() - connection.close() - second_delete = _run(db, media, "--cleanup", "--apply", "--offline") - assert second_delete.returncode == 0, second_delete.stderr - assert (media / f"{ref.media_id}.bin").exists() - - connection = sqlite3.connect(db) - connection.execute("DELETE FROM conversations WHERE inner_conversation_id=2") - connection.commit() - connection.close() - last_delete = _run(db, media, "--cleanup", "--apply", "--offline") - assert last_delete.returncode == 0, last_delete.stderr - assert not (media / f"{ref.media_id}.bin").exists() - quarantine = media.parent / "media.quarantine" - assert (quarantine / f"{ref.media_id}.bin").exists() - - -def test_cleanup_leaves_unpaired_and_symlink_objects(tmp_path): - db = tmp_path / "db.sqlite" - media = tmp_path / "media" - _db(db, []) - media.mkdir() - media_id = "b" * 64 - (media / f"{media_id}.bin").write_bytes(b"bad") - (media / f"{media_id}.json").write_text("{}") - (media / ("c" * 64 + ".bin")).write_bytes(b"bad") - result = _run(db, media, "--cleanup", "--apply", "--offline") - assert result.returncode == 0 - assert (media / f"{media_id}.bin").exists() - - -def test_restore_refuses_newer_history_and_wal_is_not_left_stale(tmp_path): - data = _image() - history = [ - { - "role": "user", - "content": [ - { - "type": "image_url", - "image_url": { - "url": "data:image/png;base64," - + base64.b64encode(data).decode() - }, - } - ], - } - ] - db = tmp_path / "db.sqlite" - media = tmp_path / "media" - backup = tmp_path / "backup" - _db(db, history) - assert ( - _run(db, media, "--apply", "--offline", "--backup", str(backup)).returncode == 0 - ) - connection = sqlite3.connect(db) - connection.execute("UPDATE conversations SET content='[]'") - connection.execute("CREATE TABLE unrelated (value TEXT)") - connection.execute("INSERT INTO unrelated VALUES ('changed')") - connection.commit() - connection.close() - (db.with_name(db.name + "-wal")).write_bytes(b"stale") - result = _run(db, media, "--restore", str(backup), "--offline") - assert result.returncode != 0 - assert not (db.with_name(db.name + ".before-restore")).exists() - - -def test_restore_validates_backup_and_conflicting_media_before_database_write(tmp_path): - db = tmp_path / "db.sqlite" - media = tmp_path / "media" - backup = tmp_path / "backup" - _db(db, []) - assert ( - _run(db, media, "--apply", "--offline", "--backup", str(backup)).returncode == 0 - ) - manifest = backup / "manifest.json" - manifest.write_text( - manifest.read_text().replace( - '"database":', '"database": "tampered", "ignored":' - ) - ) - before = db.read_bytes() - result = _run(db, media, "--restore", str(backup), "--offline") - assert result.returncode != 0 - assert db.read_bytes() == before - - -def test_successful_restore_preserves_inline_history_and_detail(tmp_path): - data = _image() - history = [ - { - "role": "user", - "content": [ - { - "type": "image_url", - "image_url": { - "url": "data:image/png;base64," - + base64.b64encode(data).decode(), - "detail": "high", - }, - } - ], - } - ] - db = tmp_path / "db.sqlite" - media = tmp_path / "media" - backup = tmp_path / "backup" - _db(db, history) - assert ( - _run(db, media, "--apply", "--offline", "--backup", str(backup)).returncode == 0 - ) - result = _run(db, media, "--restore", str(backup), "--offline") - assert result.returncode == 0, result.stderr - restored = json.loads( - sqlite3.connect(db).execute("SELECT content FROM conversations").fetchone()[0] - ) - assert restored == history - - -def test_backup_reads_consistent_open_wal_database(tmp_path): - db = tmp_path / "db.sqlite" - media = tmp_path / "media" - backup = tmp_path / "backup" - connection = sqlite3.connect(db) - connection.execute("PRAGMA journal_mode=WAL") - connection.execute( - "CREATE TABLE conversations (inner_conversation_id INTEGER PRIMARY KEY, conversation_id TEXT, content JSON)" - ) - connection.execute("INSERT INTO conversations VALUES (1, 'wal', '[]')") - connection.commit() - assert (db.with_name(db.name + "-wal")).exists() - connection.close() - result = _run(db, media, "--apply", "--offline", "--backup", str(backup)) - assert result.returncode == 0, result.stderr - assert ( - sqlite3.connect(backup / "db.sqlite") - .execute("SELECT conversation_id FROM conversations") - .fetchone()[0] - == "wal" - ) - - -def test_restore_replace_failure_leaves_database_and_cleans_stage( - tmp_path, monkeypatch -): - db = tmp_path / "db.sqlite" - media = tmp_path / "media" - backup = tmp_path / "backup" - _db(db, []) - from astrbot.core.utils.image_media_store import ImageMediaStore - - ImageMediaStore(media).put(_image(), detail="high") - assert ( - _run(db, media, "--apply", "--offline", "--backup", str(backup)).returncode == 0 - ) - for path in media.iterdir(): - path.unlink() - module_spec = importlib.util.spec_from_file_location( - "manage_image_history", "scripts/manage_image_history.py" - ) - module = importlib.util.module_from_spec(module_spec) - module_spec.loader.exec_module(module) - - def fail_replace(self, target): - raise OSError("injected replace failure") - - monkeypatch.setattr(module.Path, "replace", fail_replace) - args = type( - "Args", - (), - { - "offline": True, - "restore": str(backup), - "database": str(db), - "media": str(media), - }, - ) - before = db.read_bytes() - with pytest.raises(OSError, match="injected replace failure"): - module._restore(args) - assert db.read_bytes() == before - assert not list(tmp_path.glob("astrbot-restore-media-*")) From 843f17159b3a68aabdfd57057510ef92fef52a4d Mon Sep 17 00:00:00 2001 From: zenfun Date: Thu, 17 Sep 2026 19:12:50 +0800 Subject: [PATCH 7/7] fix(media): centralize image input safeguards --- astrbot/core/agent/tool_image_cache.py | 2 + astrbot/core/computer/file_read_utils.py | 8 + .../method/agent_sub_stages/image_input.py | 4 + .../method/agent_sub_stages/internal.py | 2 + astrbot/core/utils/image_media_store.py | 16 ++ astrbot/core/utils/media_utils.py | 240 ++++++++++++++++-- docs/en/dev/openapi-scopes.md | 1 + docs/public/openapi.json | 49 ++++ docs/zh/dev/openapi-scopes.md | 1 + tests/test_computer_fs_tools.py | 26 ++ tests/test_media_utils.py | 52 +++- tests/test_storage_cleaner.py | 14 + tests/unit/test_image_media_store.py | 30 +++ tests/unit/test_image_preparation_budget.py | 1 + tests/unit/test_image_source_preparation.py | 5 +- 15 files changed, 423 insertions(+), 28 deletions(-) diff --git a/astrbot/core/agent/tool_image_cache.py b/astrbot/core/agent/tool_image_cache.py index c873e2c58c..a84eddd31f 100644 --- a/astrbot/core/agent/tool_image_cache.py +++ b/astrbot/core/agent/tool_image_cache.py @@ -11,6 +11,7 @@ from astrbot import logger from astrbot.core.utils.astrbot_path import get_astrbot_temp_path +from astrbot.core.utils.media_utils import validate_image_input_size @dataclass @@ -91,6 +92,7 @@ def save_image( try: # Runtime cache cleanup may remove empty subdirectories. os.makedirs(self._cache_dir, exist_ok=True) + validate_image_input_size(base64_data) image_bytes = base64.b64decode(base64_data) with open(file_path, "wb") as f: f.write(image_bytes) diff --git a/astrbot/core/computer/file_read_utils.py b/astrbot/core/computer/file_read_utils.py index 55a48c0f6b..d2667f6a66 100644 --- a/astrbot/core/computer/file_read_utils.py +++ b/astrbot/core/computer/file_read_utils.py @@ -16,6 +16,10 @@ from astrbot.core.agent.context.token_counter import EstimateTokenCounter from astrbot.core.agent.message import Message from astrbot.core.agent.tool import ToolExecResult +from astrbot.core.utils.media_utils import ( + ImagePayloadTooLargeError, + validate_image_input_size, +) from .booters.base import ComputerBooter from .local_file_security import open_file_in_allowed_roots @@ -703,6 +707,10 @@ async def read_file_tool_result( return "Error reading file: binary files are not supported by this tool." if probe.kind == "image": + try: + validate_image_input_size(size_bytes) + except ImagePayloadTooLargeError as exc: + return f"Error reading file: {exc}" if local_mode: try: raw_bytes = await _read_local_file_bytes(path, local_file_descriptor) diff --git a/astrbot/core/pipeline/process_stage/method/agent_sub_stages/image_input.py b/astrbot/core/pipeline/process_stage/method/agent_sub_stages/image_input.py index 61e237b8bd..1a70ce2039 100644 --- a/astrbot/core/pipeline/process_stage/method/agent_sub_stages/image_input.py +++ b/astrbot/core/pipeline/process_stage/method/agent_sub_stages/image_input.py @@ -25,6 +25,7 @@ async def prepare_request_images( prepared: dict[str, str | None], quote_image_ref: str | None = None, montage_max_size: int | None = None, + preserve_bytes: bool = False, ) -> None: """Replace current images on a working request and track their owned files. @@ -38,6 +39,8 @@ async def prepare_request_images( prepared: Per-request mapping reused after the request hook. quote_image_ref: Optional input for the dedicated quote caption branch. montage_max_size: Optional montage-specific limit; defaults to ``max_size``. + preserve_bytes: Preserve supported still-image bytes for coordinate-sensitive + consumers. """ req.image_urls = normalize_and_dedupe_strings(req.image_urls) refs = list(req.image_urls) @@ -59,6 +62,7 @@ async def prepare_request_images( output_dir=output_dir, quality=quality, montage_max_size=montage_max_size, + preserve_bytes=preserve_bytes, ) if path: event.track_temporary_local_file(path) diff --git a/astrbot/core/pipeline/process_stage/method/agent_sub_stages/internal.py b/astrbot/core/pipeline/process_stage/method/agent_sub_stages/internal.py index 21bf1a49b0..e9272b7dca 100644 --- a/astrbot/core/pipeline/process_stage/method/agent_sub_stages/internal.py +++ b/astrbot/core/pipeline/process_stage/method/agent_sub_stages/internal.py @@ -294,6 +294,7 @@ async def process( prepared=prepared, quote_image_ref=quote_image_ref, montage_max_size=montage_max_size, + preserve_bytes=cua_pixel_mode, ) await _process_quote_message( event, @@ -361,6 +362,7 @@ async def process( output_dir=output_dir, prepared=prepared, montage_max_size=montage_max_size, + preserve_bytes=cua_pixel_mode, ) if cua_pixel_mode: oversized = [] diff --git a/astrbot/core/utils/image_media_store.py b/astrbot/core/utils/image_media_store.py index 552a3496c1..e33c4ac46d 100644 --- a/astrbot/core/utils/image_media_store.py +++ b/astrbot/core/utils/image_media_store.py @@ -13,6 +13,8 @@ from PIL import Image +from astrbot.core.utils.media_utils import validate_image_input_size + @dataclass(frozen=True, slots=True) class ImageMediaRef: @@ -253,6 +255,20 @@ def persist_inline_image_refs( """ import copy + for message in history: + parts = message.get("content") if isinstance(message, dict) else None + if not isinstance(parts, list): + continue + for part in parts: + if not isinstance(part, dict) or part.get("type") != "image_url": + continue + if part.get("_no_save"): + continue + image_url = part.get("image_url") + url = image_url.get("url") if isinstance(image_url, dict) else None + if isinstance(url, str) and url.startswith("data:image/"): + validate_image_input_size(url) + result = copy.deepcopy(history) for message in result: parts = message.get("content") if isinstance(message, dict) else None diff --git a/astrbot/core/utils/media_utils.py b/astrbot/core/utils/media_utils.py index ce25658132..cfbd43841a 100644 --- a/astrbot/core/utils/media_utils.py +++ b/astrbot/core/utils/media_utils.py @@ -43,11 +43,12 @@ 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. +# Image inputs larger than this are rejected before decoding whenever their +# encoded size is known. 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 +IMAGE_COMPRESS_DEFAULT_MIN_FILE_SIZE_BYTES = int( + IMAGE_COMPRESS_DEFAULT_MIN_FILE_SIZE_MB * 1024 * 1024 +) IMAGE_COMPRESS_DEFAULT_MAX_ENCODED_BYTES = 4 * 1024 * 1024 _WEBP_PRESERVE = object() @@ -356,6 +357,7 @@ def _open_static_webp( max_encoded_bytes: int, *, preserve_dimensions: bool = False, + preserve_input: bool = True, ): """Decode static WebP directly into one caller-owned pixel buffer. @@ -365,6 +367,8 @@ def _open_static_webp( max_encoded_bytes: Base64 payload budget. preserve_dimensions: Whether the eventual preparation step must retain the original dimensions. + preserve_input: Whether a compliant static image may be returned + without re-encoding. Returns: A Pillow image backed by one RGB/RGBA buffer, ``_WEBP_PRESERVE`` for a @@ -391,8 +395,10 @@ def _open_static_webp( f"Animated image exceeds the {max_encoded_bytes}-byte encoding limit" ) return _WEBP_PRESERVE - if encoded_size <= max_encoded_bytes and ( - preserve_dimensions or max(width, height) <= max_size + if ( + preserve_input + and encoded_size <= max_encoded_bytes + and (preserve_dimensions or max(width, height) <= max_size) ): return _WEBP_PRESERVE # The direct route cannot carry EXIF orientation without loading Pillow's @@ -800,6 +806,124 @@ def _decode_base64_payload( raise ValueError(error_message) from exc +def _estimate_base64_decoded_size( + payload: str, + *, + start: int = 0, + require_valid_chars: bool = False, +) -> int | None: + """Estimate decoded bytes without creating a compact payload copy. + + Args: + payload: Base64 text, possibly containing whitespace. + start: Index at which the base64 payload begins. + require_valid_chars: Whether to reject characters outside standard + base64 and whitespace. + + Returns: + The estimated decoded byte count, or ``None`` when the text is not a + complete standard base64 payload. + """ + encoded_size = 0 + padding_size = 0 + for index, char in enumerate(payload): + if index < start: + continue + if char.isspace(): + continue + if require_valid_chars and char not in ( + "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/=" + ): + return None + encoded_size += 1 + if char == "=": + padding_size += 1 + + if encoded_size % 4 == 1: + return None + return max(0, encoded_size * 3 // 4 - padding_size) + + +def validate_image_input_size(image_source: bytes | str | Path | int) -> int | None: + """Reject an image source that is too large to decode safely. + + Args: + image_source: Raw bytes, a local path, an image reference, or a known + byte count. HTTP(S) references return ``None`` because their size + is unknown until download completes. + + Returns: + The known encoded byte count, or ``None`` when the source size cannot + be determined without reading it. + + Raises: + ImagePayloadTooLargeError: The known input exceeds the image input cap. + TypeError: The source type is unsupported. + OSError: A local file cannot be inspected. + """ + source_size: int | None = None + if isinstance(image_source, bool): + raise TypeError("Image source size must not be a boolean") + if isinstance(image_source, int): + source_size = image_source + elif isinstance(image_source, bytes): + source_size = len(image_source) + elif isinstance(image_source, Path): + try: + source_size = image_source.stat().st_size + except FileNotFoundError: + return None + elif isinstance(image_source, str): + if image_source.startswith(("http://", "https://")): + return None + if image_source.startswith("data:"): + comma_index = image_source.find(",") + if comma_index < 0: + return None + header = image_source[:comma_index] + header_parts = header[5:].split(";") + if any(part.lower() == "base64" for part in header_parts[1:]): + source_size = _estimate_base64_decoded_size( + image_source, + start=comma_index + 1, + ) + elif image_source.startswith("base64://"): + source_size = _estimate_base64_decoded_size( + image_source, + start=len("base64://"), + ) + else: + is_uri = is_file_uri(image_source) + if is_uri: + try: + source_size = Path(file_uri_to_path(image_source)).stat().st_size + except FileNotFoundError: + source_size = None + else: + try: + source_size = Path(image_source).stat().st_size + except (FileNotFoundError, OSError, ValueError): + source_size = None + if source_size is None and not is_uri: + source_size = _estimate_base64_decoded_size( + image_source, + require_valid_chars=True, + ) + else: + raise TypeError(f"Unsupported image source type: {type(image_source).__name__}") + + if source_size is None: + return None + if source_size < 0: + raise ValueError("Image source size must not be negative") + if source_size > MODEL_IMAGE_MAX_INPUT_BYTES: + raise ImagePayloadTooLargeError( + "Image input exceeds the " + f"{MODEL_IMAGE_MAX_INPUT_BYTES}-byte limit ({source_size} bytes)" + ) + return source_size + + def _encode_file_to_base64(path: Path) -> str: """Encode a file without retaining a second full raw-image copy. @@ -995,6 +1119,8 @@ async def _materialize_media_ref( cleanup_paths.append(target_path) try: await download_file(media_ref, str(target_path)) + if media_type == "image": + validate_image_input_size(target_path) except Exception: _cleanup_paths(cleanup_paths) raise @@ -1020,9 +1146,13 @@ async def _materialize_media_ref( if is_file_uri(media_ref): path = Path(file_uri_to_path(media_ref)) + if media_type == "image": + validate_image_input_size(path) return _LocalMediaFile(path=path, mime_type=_guess_mime_type(path)) if media_ref.startswith("data:"): + if media_type == "image": + validate_image_input_size(media_ref) mime_type, media_bytes = _parse_base64_data_uri(media_ref) target_suffix = _extension_from_mime_type(mime_type) or suffix if media_type == "image" and target_suffix == suffix: @@ -1047,6 +1177,8 @@ async def _materialize_media_ref( ) if media_ref.startswith("base64://"): + if media_type == "image": + validate_image_input_size(media_ref) media_bytes = _decode_base64_payload( media_ref.removeprefix("base64://"), error_message="invalid base64 media payload", @@ -1079,8 +1211,12 @@ async def _materialize_media_ref( except OSError: pass if path_exists: + if media_type == "image": + validate_image_input_size(path) return _LocalMediaFile(path=path, mime_type=_guess_mime_type(path)) + if media_type == "image": + validate_image_input_size(media_ref) compact_media_ref = "".join(media_ref.split()) if compact_media_ref: try: @@ -1795,7 +1931,11 @@ def _publish_image_cache_atomic( def _convert_image_bytes_sync( - source_bytes: bytes, max_size: int, quality: int + source_bytes: bytes, + max_size: int, + quality: int, + *, + preserve_bytes: bool = False, ) -> bytes: """Normalize a validated still image with an optional derived cache. @@ -1803,17 +1943,23 @@ def _convert_image_bytes_sync( source_bytes: Encoded source bytes already checked by _inspect_image. max_size: Longest-edge limit in pixels. quality: JPEG output quality in the range 1-100. + preserve_bytes: Keep a supported, correctly oriented still byte-exact + even when it exceeds the normal small-source threshold. Returns: Single-frame JPEG or PNG bytes. An oriented JPEG or PNG within the size - and reuse-byte limits is reused unchanged; anything else is re-encoded. + and small-source limit is reused unchanged; anything else is re-encoded. """ + validate_image_input_size(source_bytes) 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 + and ( + preserve_bytes + or len(source_bytes) <= IMAGE_COMPRESS_DEFAULT_MIN_FILE_SIZE_BYTES + ) ): return source_bytes cache_key = _image_convert_cache_key( @@ -1939,6 +2085,7 @@ async def prepare_model_image( output_dir: Path, quality: int = IMAGE_COMPRESS_DEFAULT_QUALITY, montage_max_size: int | None = None, + preserve_bytes: bool = False, ) -> str | None: """Prepare a single local model-ready image for the caller to own until consumption. @@ -1952,6 +2099,8 @@ async def prepare_model_image( but montages are never used for coordinates, so callers pass the configured limit here to keep the 3x3 canvas bounded. Defaults to ``max_size``. + preserve_bytes: Preserve supported, correctly oriented still-image bytes + for coordinate-sensitive consumers such as CUA. Returns: An existing JPEG or PNG path, or None for a recoverable input or write @@ -1960,14 +2109,6 @@ 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) if frame_count > 1: @@ -1979,7 +2120,11 @@ async def prepare_model_image( ) else: converted_bytes = await asyncio.to_thread( - _convert_image_bytes_sync, image_bytes, max_size, quality + _convert_image_bytes_sync, + image_bytes, + max_size, + quality, + preserve_bytes=preserve_bytes, ) # Publish the working file synchronously after encoding, so cancellation # cannot leave an untracked background write alive after this call. @@ -2674,7 +2819,12 @@ def _compress_image_sync( """ if max_size < 1 or max_encoded_bytes < 1 or not 1 <= quality <= 100: raise ValueError("Image dimensions, byte budget and quality must be positive") - source_bytes = len(source) if isinstance(source, bytes) else source.stat().st_size + if isinstance(source, bytes): + source_bytes = len(source) + else: + source_bytes = validate_image_input_size(source) + if source_bytes is None: + source_bytes = source.stat().st_size encoded_size = 4 * ((source_bytes + 2) // 3) direct_webp = None if isinstance(source, bytes): @@ -2691,6 +2841,7 @@ def _compress_image_sync( max_size, max_encoded_bytes, preserve_dimensions=preserve_dimensions, + preserve_input=source_bytes <= IMAGE_COMPRESS_DEFAULT_MIN_FILE_SIZE_BYTES, ) if direct_webp is _WEBP_PRESERVE: return None @@ -2705,8 +2856,10 @@ def _compress_image_sync( f"Animated image exceeds the {max_encoded_bytes}-byte encoding limit" ) return None - if encoded_size <= max_encoded_bytes and ( - preserve_dimensions or max(opened.size) <= max_size + if ( + source_bytes <= IMAGE_COMPRESS_DEFAULT_MIN_FILE_SIZE_BYTES + and encoded_size <= max_encoded_bytes + and (preserve_dimensions or max(opened.size) <= max_size) ): return None @@ -2877,25 +3030,52 @@ async def compress_image( """ if url_or_path.startswith(("http://", "https://")): return url_or_path + max_size = max(int(max_size), 1) + quality = min(max(int(quality), 1), 100) + max_encoded_bytes = max(int(max_encoded_bytes), 1) + image_source: bytes | Path + source_size: int + + def _fits_max_size(source: bytes | Path) -> bool | None: + try: + image_file = io.BytesIO(source) if isinstance(source, bytes) else source + with PILImage.open(image_file) as opened_image: + return max(opened_image.size) <= max_size + except MemoryError: + raise + except Exception: + return None + if url_or_path.startswith("data:image"): + validate_image_input_size(url_or_path) _, encoded = url_or_path.split(",", 1) image_source: bytes | Path = _decode_base64_payload( encoded, error_message="invalid image data URI payload" ) + source_size = len(image_source) else: image_source = Path(url_or_path) - if not image_source.exists(): + source_size = validate_image_input_size(image_source) + if source_size is None: return url_or_path + encoded_size = 4 * ((source_size + 2) // 3) + if ( + source_size < IMAGE_COMPRESS_DEFAULT_MIN_FILE_SIZE_BYTES + and encoded_size <= max_encoded_bytes + and _fits_max_size(image_source) + ): + return url_or_path + worker = asyncio.create_task( asyncio.to_thread( _compress_image_sync, image_source, Path(get_astrbot_temp_path()), - max(int(max_size), 1), - min(max(int(quality), 1), 100), + max_size, + quality, optimize, - max(int(max_encoded_bytes), 1), + max_encoded_bytes, preserve_dimensions=preserve_dimensions, ) ) @@ -2952,6 +3132,7 @@ async def _prepare() -> ResolvedMediaData: try: resolved_source = source_ref if isinstance(source_ref, bytes): + validate_image_input_size(source_ref) owned_source = _temp_media_path("image", ".bin") await asyncio.to_thread(owned_source.write_bytes, source_ref) resolved_source = str(owned_source) @@ -3011,8 +3192,15 @@ async def _prepare() -> ResolvedMediaData: finally: if output_path != resolved.path: output_path.unlink(missing_ok=True) + if mime_type == "application/octet-stream": + mime_type = None + if mime_type is None and resolved.mime_type not in { + None, + "application/octet-stream", + }: + mime_type = resolved.mime_type if not mime_type: - mime_type = resolved.mime_type or default_mime_type + mime_type = default_mime_type if not mime_type: raise ValueError( f"Invalid image file: {describe_media_ref(resolved_source)}" diff --git a/docs/en/dev/openapi-scopes.md b/docs/en/dev/openapi-scopes.md index 4bf87c6c4b..7b44c6e9b7 100644 --- a/docs/en/dev/openapi-scopes.md +++ b/docs/en/dev/openapi-scopes.md @@ -184,6 +184,7 @@ Manage conversations and platform-session data. | `GET` | `/api/v1/conversations/{conversation_id}` | — | | `PATCH` | `/api/v1/conversations/{conversation_id}` | — | | `DELETE` | `/api/v1/conversations/{conversation_id}` | — | +| `GET` | `/api/v1/conversations/{conversation_id}/media/{media_id}` | — | | `PUT` | `/api/v1/conversations/{conversation_id}/messages` | — | | `GET` | `/api/v1/session-groups` | — | | `POST` | `/api/v1/session-groups` | — | diff --git a/docs/public/openapi.json b/docs/public/openapi.json index b666e92afa..512eca8393 100644 --- a/docs/public/openapi.json +++ b/docs/public/openapi.json @@ -6418,6 +6418,55 @@ "description": "**Required scope:** `data`" } }, + "/api/v1/conversations/{conversation_id}/media/{media_id}": { + "get": { + "tags": [ + "Conversations" + ], + "summary": "Preview an image referenced by a conversation", + "operationId": "previewConversationMedia", + "x-astrbot-scope": "data", + "parameters": [ + { + "$ref": "#/components/parameters/ConversationId" + }, + { + "name": "media_id", + "in": "path", + "required": true, + "schema": { + "type": "string", + "pattern": "^[0-9a-f]{64}$" + } + }, + { + "name": "user_id", + "in": "query", + "required": true, + "schema": { + "type": "string" + } + } + ], + "responses": { + "200": { + "description": "Image bytes", + "content": { + "image/*": { + "schema": { + "type": "string", + "format": "binary" + } + } + } + }, + "404": { + "description": "Conversation or media is unavailable" + } + }, + "description": "**Required scope:** `data`" + } + }, "/api/v1/conversations/{conversation_id}/messages": { "put": { "tags": [ diff --git a/docs/zh/dev/openapi-scopes.md b/docs/zh/dev/openapi-scopes.md index 3152454eb3..36146c7ef8 100644 --- a/docs/zh/dev/openapi-scopes.md +++ b/docs/zh/dev/openapi-scopes.md @@ -184,6 +184,7 @@ outline: deep | `GET` | `/api/v1/conversations/{conversation_id}` | — | | `PATCH` | `/api/v1/conversations/{conversation_id}` | — | | `DELETE` | `/api/v1/conversations/{conversation_id}` | — | +| `GET` | `/api/v1/conversations/{conversation_id}/media/{media_id}` | — | | `PUT` | `/api/v1/conversations/{conversation_id}/messages` | — | | `GET` | `/api/v1/session-groups` | — | | `POST` | `/api/v1/session-groups` | — | diff --git a/tests/test_computer_fs_tools.py b/tests/test_computer_fs_tools.py index 7e224e1bab..d5b91d892f 100644 --- a/tests/test_computer_fs_tools.py +++ b/tests/test_computer_fs_tools.py @@ -19,6 +19,7 @@ from astrbot.core.computer.booters.local import LocalBooter from astrbot.core.tools.computer_tools import fs as fs_tools from astrbot.core.tools.computer_tools import util as computer_util +from astrbot.core.utils import media_utils def _make_context( @@ -859,6 +860,31 @@ async def test_file_read_tool_returns_image_call_tool_result_for_images( assert base64.b64decode(result.content[0].data) == image_path.read_bytes() +@pytest.mark.asyncio +@pytest.mark.skipif(os.name == "nt", reason="Restricted file access needs POSIX.") +async def test_file_read_tool_rejects_oversized_image_before_reading( + monkeypatch: pytest.MonkeyPatch, + tmp_path, +): + workspace = _setup_local_fs_tools(monkeypatch, tmp_path) + image_path = workspace / "oversized.png" + image_path.write_bytes(b"\x89PNG\r\n\x1a\n") + with image_path.open("ab") as image_file: + image_file.truncate(media_utils.MODEL_IMAGE_MAX_INPUT_BYTES + 1) + + async def fail_read(*_args, **_kwargs): + raise AssertionError("oversized image must be rejected before reading") + + monkeypatch.setattr(file_read_utils, "_read_local_file_bytes", fail_read) + result = await fs_tools.FileReadTool().call( + _make_context(), + path="oversized.png", + ) + + assert isinstance(result, str) + assert "Image input exceeds" in result + + @pytest.mark.asyncio async def test_local_file_read_defers_image_preparation_to_tool_loop( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/test_media_utils.py b/tests/test_media_utils.py index 08402cfdde..05f8e3aa78 100644 --- a/tests/test_media_utils.py +++ b/tests/test_media_utils.py @@ -324,6 +324,34 @@ async def test_compress_image_preserves_alpha_png(tmp_path, monkeypatch): compressed_path.unlink(missing_ok=True) +@pytest.mark.asyncio +async def test_compress_image_reencodes_sources_over_one_megabyte( + tmp_path, monkeypatch +): + """Sources above the legacy threshold still enter the compression path.""" + from PIL import Image as PILImage + + temp_dir = tmp_path / "temp" + monkeypatch.setattr(media_utils, "get_astrbot_temp_path", lambda: str(temp_dir)) + image_path = tmp_path / "large.png" + with PILImage.frombytes("RGB", (768, 768), os.urandom(768 * 768 * 3)) as image: + image.save(image_path, format="PNG") + + assert ( + image_path.stat().st_size + > media_utils.IMAGE_COMPRESS_DEFAULT_MIN_FILE_SIZE_BYTES + ) + compressed_path = Path(await media_utils.compress_image(str(image_path))) + + try: + assert compressed_path != image_path + assert compressed_path.suffix == ".jpg" + assert compressed_path.read_bytes().startswith(b"\xff\xd8") + assert image_path.read_bytes().startswith(b"\x89PNG") + finally: + compressed_path.unlink(missing_ok=True) + + @pytest.mark.asyncio async def test_compress_image_rejects_oversized_encoded_payload(tmp_path, monkeypatch): from PIL import Image as PILImage @@ -910,6 +938,11 @@ async def test_prepare_model_image_skips_oversized_input(tmp_path, monkeypatch): with image_path.open("ab") as f: f.truncate(media_utils.MODEL_IMAGE_MAX_INPUT_BYTES + 1) + def fail_read(*_args, **_kwargs): + raise AssertionError("oversized image must be rejected before reading") + + monkeypatch.setattr(Path, "read_bytes", fail_read) + result = await media_utils.prepare_model_image( str(image_path), max_size=1280, output_dir=tmp_path ) @@ -917,6 +950,23 @@ async def test_prepare_model_image_skips_oversized_input(tmp_path, monkeypatch): assert result is None +@pytest.mark.asyncio +async def test_prepare_image_source_rejects_oversized_data_uri_before_decode( + monkeypatch, +): + """Inline payloads must be bounded before base64 decoding allocates bytes.""" + monkeypatch.setattr(media_utils, "MODEL_IMAGE_MAX_INPUT_BYTES", 4) + + def fail_decode(*_args, **_kwargs): + raise AssertionError("oversized data URI must not be decoded") + + monkeypatch.setattr(media_utils, "_decode_base64_payload", fail_decode) + image_ref = "data:image/png;base64," + base64.b64encode(b"12345").decode() + + with pytest.raises(media_utils.ImagePayloadTooLargeError, match="input exceeds"): + await media_utils.prepare_image_source(image_ref) + + 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 @@ -940,7 +990,7 @@ def test_convert_image_bytes_reencodes_large_in_range_input(tmp_path, monkeypatc buffer = BytesIO() img.save(buffer, format="PNG") source = buffer.getvalue() - assert len(source) > media_utils.MODEL_IMAGE_REUSE_MAX_BYTES + assert len(source) > media_utils.IMAGE_COMPRESS_DEFAULT_MIN_FILE_SIZE_BYTES result = media_utils._convert_image_bytes_sync(source, 1280, 95) diff --git a/tests/test_storage_cleaner.py b/tests/test_storage_cleaner.py index 5d6f82c515..719d625636 100644 --- a/tests/test_storage_cleaner.py +++ b/tests/test_storage_cleaner.py @@ -1,7 +1,10 @@ import base64 from pathlib import Path +import pytest + from astrbot.core.agent.tool_image_cache import tool_image_cache +from astrbot.core.utils import media_utils from astrbot.core.utils.storage_cleaner import StorageCleaner @@ -111,3 +114,14 @@ def test_tool_image_cache_recovers_after_storage_cleanup(tmp_path, monkeypatch): assert cached_path == cache_dir / "call-test_0.jpg" assert cached_path.read_bytes() == image_bytes + + +def test_tool_image_cache_rejects_oversized_image_before_decode(monkeypatch): + monkeypatch.setattr(media_utils, "MODEL_IMAGE_MAX_INPUT_BYTES", 4) + + with pytest.raises(media_utils.ImagePayloadTooLargeError, match="input exceeds"): + tool_image_cache.save_image( + base64_data=base64.b64encode(b"12345").decode("ascii"), + tool_call_id="call-too-large", + tool_name="test-tool", + ) diff --git a/tests/unit/test_image_media_store.py b/tests/unit/test_image_media_store.py index 4adc030fe9..937987c1ae 100644 --- a/tests/unit/test_image_media_store.py +++ b/tests/unit/test_image_media_store.py @@ -1,5 +1,6 @@ """Tests for durable image object ownership and authorization.""" +import base64 import io import json import os @@ -8,6 +9,7 @@ import pytest from PIL import Image +from astrbot.core.utils import media_utils from astrbot.core.utils.image_media_store import ImageMediaStore @@ -67,6 +69,34 @@ def test_store_never_commits_invalid_or_partial_media(tmp_path): assert not list((tmp_path / "media").glob("*.json")) +def test_persist_inline_image_refs_rejects_oversized_payload_before_decode( + tmp_path, monkeypatch +): + store = ImageMediaStore(tmp_path / "media") + monkeypatch.setattr(media_utils, "MODEL_IMAGE_MAX_INPUT_BYTES", 4) + encoded = base64.b64encode(b"12345").decode("ascii") + + with pytest.raises(media_utils.ImagePayloadTooLargeError, match="input exceeds"): + from astrbot.core.utils.image_media_store import persist_inline_image_refs + + persist_inline_image_refs( + [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": f"data:image/png;base64,{encoded}"}, + } + ], + } + ], + store, + ) + + assert not list((tmp_path / "media").glob("*")) + + @pytest.mark.asyncio @pytest.mark.parametrize("runtime_message", [False, True]) async def test_reference_materialization_preserves_history_and_bytes( diff --git a/tests/unit/test_image_preparation_budget.py b/tests/unit/test_image_preparation_budget.py index 2abc322114..86aafe219e 100644 --- a/tests/unit/test_image_preparation_budget.py +++ b/tests/unit/test_image_preparation_budget.py @@ -338,6 +338,7 @@ async def fail(*args, **kwargs): async def test_cancelled_worker_cleans_after_exit(tmp_path, monkeypatch, timeout): source = tmp_path / "source" source.write_bytes(b"source") + monkeypatch.setattr(media_utils, "IMAGE_COMPRESS_DEFAULT_MIN_FILE_SIZE_BYTES", 1) output = tmp_path / "worker-output" entered = threading.Event() release = threading.Event() diff --git a/tests/unit/test_image_source_preparation.py b/tests/unit/test_image_source_preparation.py index 50c46a1c61..70de889f46 100644 --- a/tests/unit/test_image_source_preparation.py +++ b/tests/unit/test_image_source_preparation.py @@ -70,7 +70,10 @@ def blocked(source, *args, **kwargs): monkeypatch.setattr(media_utils, "_compress_image_sync", blocked) encoded = base64.b64encode(_png()).decode() task = asyncio.create_task( - media_utils.prepare_image_source(f"data:image/png;base64,{encoded}") + media_utils.prepare_image_source( + f"data:image/png;base64,{encoded}", + options=media_utils.ImagePreparationOptions(max_size=2), + ) ) assert await asyncio.to_thread(entered.wait, 2) assert source_seen["path"] is not None