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/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/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..d2667f6a66 100644 --- a/astrbot/core/computer/file_read_utils.py +++ b/astrbot/core/computer/file_read_utils.py @@ -16,12 +16,9 @@ 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, + ImagePayloadTooLargeError, + validate_image_input_size, ) from .booters.base import ComputerBooter @@ -309,25 +306,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 +321,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" @@ -756,33 +707,34 @@ 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: - 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/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 322506e3a8..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 = [] @@ -538,12 +540,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 15fd6b72f4..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 ] @@ -371,6 +420,9 @@ def __init__(self, provider_config, provider_settings) -> None: default_headers=self.custom_headers, base_url=provider_config.get("api_base") or None, timeout=self.timeout, + # Retry is handled by retry_provider_request(); disable the + # SDK built-in retry to avoid stacking request attempts. + max_retries=0, http_client=self._create_http_client(provider_config), ) else: @@ -380,6 +432,9 @@ def __init__(self, provider_config, provider_settings) -> None: base_url=provider_config.get("api_base") or None, default_headers=self.custom_headers, timeout=self.timeout, + # Retry is handled by retry_provider_request(); disable the + # SDK built-in retry to avoid stacking request attempts. + max_retries=0, http_client=self._create_http_client(provider_config), ) @@ -955,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( @@ -1075,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..e33c4ac46d --- /dev/null +++ b/astrbot/core/utils/image_media_store.py @@ -0,0 +1,373 @@ +"""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 + +from astrbot.core.utils.media_utils import validate_image_input_size + + +@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 + + 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 + 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 ff552293d5..cfbd43841a 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 @@ -42,6 +43,472 @@ IMAGE_COMPRESS_DEFAULT_QUALITY = 95 IMAGE_COMPRESS_DEFAULT_OPTIMIZE = True IMAGE_COMPRESS_DEFAULT_MIN_FILE_SIZE_MB = 1.0 +# Image inputs larger than this are rejected before decoding whenever their +# encoded size is known. +MODEL_IMAGE_MAX_INPUT_BYTES = 32 * 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() +_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, + preserve_input: bool = True, +): + """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. + 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 + 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 ( + 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 + # 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", @@ -120,6 +587,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. @@ -133,6 +616,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.""" @@ -175,7 +659,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.""" @@ -322,6 +806,157 @@ 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. + + 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. @@ -484,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 @@ -509,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: @@ -536,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", @@ -568,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: @@ -911,29 +1558,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. + + Returns: + Prepared provider-ready image data, or ``None`` for a safely skippable + invalid image when ``strict`` is false. - ``strict=False`` returns ``None`` for invalid images so payload - assembly can skip bad image refs without failing the whole request. + 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: @@ -1224,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. @@ -1232,16 +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 - limit 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 ( + preserve_bytes + or len(source_bytes) <= IMAGE_COMPRESS_DEFAULT_MIN_FILE_SIZE_BYTES + ) ): return source_bytes cache_key = _image_convert_cache_key( @@ -1367,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. @@ -1380,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 @@ -1399,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. @@ -1453,10 +2178,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. @@ -1465,7 +2192,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) @@ -2009,134 +2741,487 @@ 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") + 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): + 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, + preserve_input=source_bytes <= IMAGE_COMPRESS_DEFAULT_MIN_FILE_SIZE_BYTES, + ) + 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 ( + 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 + 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. + The original reference or a caller-owned prepared file path. + + 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 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 + max_encoded_bytes = max(int(max_encoded_bytes), 1) + image_source: bytes | Path + source_size: int - def _exceeds_max_size(source: bytes | Path) -> bool: + def _fits_max_size(source: bytes | Path) -> bool | None: 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"): - 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", + 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" ) - if len(image_source) < min_file_size_bytes and not _exceeds_max_size( - image_source - ): - return url_or_path + source_size = len(image_source) else: - local_path = Path(url_or_path) - if not local_path.exists(): + image_source = Path(url_or_path) + source_size = validate_image_input_size(image_source) + if source_size is None: 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: + 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 - temp_dir = Path(get_astrbot_temp_path()) - temp_dir.mkdir(parents=True, exist_ok=True) - - # 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, + worker = asyncio.create_task( + asyncio.to_thread( + _compress_image_sync, + image_source, + Path(get_astrbot_temp_path()), + max_size, + quality, + optimize, + max_encoded_bytes, + 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") + + worker.add_done_callback(cleanup_finished) + raise return compressed_path or url_or_path + + +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) + ) + source_ref = preparation_input.value + + async def _prepare() -> ResolvedMediaData: + owned_source: Path | None = None + 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) + 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 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 = 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/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/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/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/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..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/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/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_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_computer_fs_tools.py b/tests/test_computer_fs_tools.py index a2f26670d8..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( @@ -184,12 +185,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 +856,72 @@ 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 +@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, + 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 976efc6db3..05f8e3aa78 100644 --- a/tests/test_media_utils.py +++ b/tests/test_media_utils.py @@ -324,6 +324,58 @@ 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 + + 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 +395,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 @@ -873,3 +925,75 @@ 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) + + 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 + ) + + 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 + + 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.IMAGE_COMPRESS_DEFAULT_MIN_FILE_SIZE_BYTES + + result = media_utils._convert_image_bytes_sync(source, 1280, 95) + + assert result is not source + assert result[:2] == b"\xff\xd8" + assert len(result) < len(source) diff --git a/tests/test_openai_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 e92d90f8d8..1fe68e317c 100644 --- a/tests/test_openai_source.py +++ b/tests/test_openai_source.py @@ -1,10 +1,12 @@ import base64 import builtins from io import BytesIO +from pathlib import Path from types import SimpleNamespace import httpx import pytest +from openai import AsyncAzureOpenAI, AsyncOpenAI from openai.types.chat.chat_completion import ChatCompletion from openai.types.chat.chat_completion_chunk import ChatCompletionChunk from PIL import Image as PILImage @@ -15,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 @@ -62,6 +62,54 @@ 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_provider_client_disables_sdk_builtin_retries(overrides, expected_client): + provider = _make_provider(overrides) + try: + assert isinstance(provider.client, expected_client) + assert provider.client.max_retries == 0 + finally: + await provider.terminate() + + +@pytest.mark.asyncio +async def test_query_attempts_exactly_request_max_retries_times(monkeypatch): + monkeypatch.setattr(request_retry, "REQUEST_RETRY_WAIT_MIN_S", 0) + monkeypatch.setattr(request_retry, "REQUEST_RETRY_WAIT_MAX_S", 0) + + provider = _make_provider() + try: + calls = 0 + + async def failing_create(**kwargs): + nonlocal calls + calls += 1 + raise httpx.ConnectError("temporary connection failure") + + monkeypatch.setattr(provider.client.chat.completions, "create", failing_create) + + with pytest.raises(httpx.ConnectError): + await provider._query( + payloads={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}], + }, + tools=None, + request_max_retries=2, + ) + + assert calls == 2 + finally: + await provider.terminate() + + def test_create_http_client_uses_openai_httpx_module(monkeypatch): captured: dict[str, object] = {} fake_httpx_module = object() @@ -725,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( @@ -995,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": { @@ -1028,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_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/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_media_store.py b/tests/unit/test_image_media_store.py new file mode 100644 index 0000000000..937987c1ae --- /dev/null +++ b/tests/unit/test_image_media_store.py @@ -0,0 +1,334 @@ +"""Tests for durable image object ownership and authorization.""" + +import base64 +import io +import json +import os +from concurrent.futures import ThreadPoolExecutor + +import pytest +from PIL import Image + +from astrbot.core.utils import media_utils +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")) + + +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( + 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_preparation_budget.py b/tests/unit/test_image_preparation_budget.py new file mode 100644 index 0000000000..86aafe219e --- /dev/null +++ b/tests/unit/test_image_preparation_budget.py @@ -0,0 +1,377 @@ +"""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") + monkeypatch.setattr(media_utils, "IMAGE_COMPRESS_DEFAULT_MIN_FILE_SIZE_BYTES", 1) + 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..70de889f46 --- /dev/null +++ b/tests/unit/test_image_source_preparation.py @@ -0,0 +1,137 @@ +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}", + options=media_utils.ImagePreparationOptions(max_size=2), + ) + ) + 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_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