From c1de31ab2cf581394b821382d78d87eeeed35fad Mon Sep 17 00:00:00 2001 From: WQS Date: Fri, 9 Oct 2026 17:42:13 +0800 Subject: [PATCH] feat(auto-memory): save session images with source links (#574) * feat(auto-memory): save session images with contextual source links * fix(auto-memory): validate image provenance before persistence * refactor(auto-memory): keep image attachment helpers in the memory step * fix(auto-memory): preserve main session persistence order * refactor(auto-memory): simplify image reuse and source frontmatter --- docs/en/auto_memory.md | 31 ++- docs/zh/auto_memory.md | 25 +- reme/config/default.yaml | 2 +- reme/steps/evolve/auto_memory.py | 161 ++++++++++++- reme/steps/evolve/auto_memory.yaml | 13 + tests/unit/test_auto_memory_direct_images.py | 148 +++++++++++- tests/unit/test_session_image_assets.py | 236 +++++++++++++++++++ 7 files changed, 597 insertions(+), 19 deletions(-) create mode 100644 tests/unit/test_session_image_assets.py diff --git a/docs/en/auto_memory.md b/docs/en/auto_memory.md index 05ba6fdb..02064837 100644 --- a/docs/en/auto_memory.md +++ b/docs/en/auto_memory.md @@ -89,9 +89,31 @@ image blocks are present, the existing text-only behavior is unchanged, includin Pass images as top-level AgentScope `DataBlock` values in `messages`, with an `image/` media type. Text and images stay in their original order, with speaker and timestamp boundaries preserved. Base64 sources and HTTP(S) URLs pass unchanged to -the formatter; Auto Memory does not download or preprocess the images. URLs must be accessible to the model provider. For local +the formatter; images are not resized or transcoded. URLs are not downloaded and must be accessible to the model provider. For local files, submit Base64 instead of a `file://` URL; other URL schemes are also unsupported. +With images enabled, Auto Memory saves each Base64 image's original bytes under the configured `session_dir`: + +```text +session/images//msg--image-. +``` + +The filename uses the message's `id` and the image's position among all content blocks, starting at zero; the extension +comes from its media type. Keep session IDs, message IDs and block positions stable when resubmitting a conversation: +an existing file at that path is reused without comparing its contents. Use a new message ID when replacing an image. +Calls with images disabled or no images do not save attachments. + +Each image is accompanied by its exact source link in the model input. The memory prompt asks the Agent to cite that source +beside the corresponding visual facts, for example: + +```markdown +The diagram places Gateway before Worker and PostgreSQL. See [[session/images/session-a/msg-6d6573736167652d61-image-1.png]]. +``` + +For URL images, the citation uses the original URL. Auto Memory also adds the supplied image sources to the daily note's +`source_images` frontmatter, preserving existing entries. This list records provenance; the body links connect individual +facts to their images. These session attachments are not watched as resources and do not trigger separate caption calls. + The wrapper's `context_config.max_image_num` limits the number of images per call; Auto Memory rejects excess images rather than increasing the limit. The AgentScope default is 5. To use a higher limit, set it when starting the service: @@ -109,9 +131,10 @@ Model and formatter limits still apply. When image input is enabled and images a backend, URL schemes and image count before saving the conversation. Later formatter or provider errors are returned without retrying as text-only. As with text-only calls, those errors do not roll back an already saved conversation. -Source JSONL saving follows the filtering rules above, including the omission of Base64 blocks. To process those images -again, resubmit the original messages rather than the saved JSONL. No separate image files or caption cards are created, -though the wrapper's internal Agent state under `mem_session/agentscope` can contain image inputs. +Source JSONL saving follows the filtering rules above, including the omission of Base64 blocks. Saved attachments are not +automatically restored into a replay of that JSONL; to process the images again, resubmit the original messages. Attachments +already saved remain available if the model call fails or decides not to write a memory card; Auto Memory does not clean +them up automatically. Local source paths can also be passed to `read_image`. ## Message Timestamps diff --git a/docs/zh/auto_memory.md b/docs/zh/auto_memory.md index 7077a1bc..7e946f15 100644 --- a/docs/zh/auto_memory.md +++ b/docs/zh/auto_memory.md @@ -80,9 +80,28 @@ Auto Memory 直接用这个模型理解图文，不先生成 caption。关闭图 wrapper 类型。 在 `messages` 中用 AgentScope 顶层 `DataBlock` 传入图像，媒体类型以 `image/` 开头。文本和图像按原顺序交错排列， -保留说话人和时间信息。Base64 source 与 HTTP(S) URL 原样交给 formatter，Auto Memory 不下载或预处理图像。URL 需要能被模型 +保留说话人和时间信息。Base64 source 与 HTTP(S) URL 原样交给 formatter，不缩放或转码。URL 不会被下载，需要能被模型 供应商访问；本地文件请先转为 Base64，不使用 `file://` URL，其他 URL scheme 也不支持。 +开启图像后，Auto Memory 会把 Base64 图像的原始字节保存到配置的 `session_dir` 下： + +```text +session/images//msg--image-. +``` + +文件名使用消息的 `id`，以及图像在所有 content block 中的位置（从零开始），扩展名取自媒体类型。重复提交对话时，请保持 +session ID、消息 ID 和 block 位置不变：同一路径已有文件会直接复用，不比较内容；更换图像时使用新的消息 ID。 +关闭图像或消息中没有图像时，不保存附件。 + +模型输入中，每张图像旁边都会带上准确的来源链接。记忆 prompt 要求 Agent 在相应的视觉事实旁引用原图，例如： + +```markdown +部署图中，Gateway 位于 Worker 和 PostgreSQL 之前，见 [[session/images/session-a/msg-6d6573736167652d61-image-1.png]]。 +``` + +URL 图像使用原始 URL 作为引用。Auto Memory 还会把本次传入的图像来源补充到 daily note 的 `source_images` frontmatter 中， +保留已有条目。这个列表负责记录来源，正文链接则说明具体事实对应哪张图。会话附件不会被当作资源监听，也不会触发额外的 caption 调用。 + 每次调用的图像数量受 wrapper 的 `context_config.max_image_num` 限制，超限会报错，不会自动提高上限。 AgentScope 默认允许 5 张图像。需要更多时，在启动服务时设置： @@ -99,8 +118,8 @@ reme auto_memory session_id=session-a include_images=true messages='[...]' 模型与 formatter 自身的限制仍然适用。开启图像且消息中包含图像时，才会在保存对话前检查 wrapper backend、URL scheme 和图像数量。 之后的 formatter 或 provider 错误直接返回，不转为纯文本重试；与纯文本调用相同，已保存的对话不会因此回滚。 -源 JSONL 仍按上文规则保存，包括过滤 Base64 block。因此，再次处理这些图像需要提交原始消息，而不是读取已保存的 JSONL。 -不会另外生成图像文件或 caption 卡片，但 wrapper 保存在 `mem_session/agentscope` 中的内部 Agent 状态可能包含图像输入。 +源 JSONL 仍按上文规则保存，包括过滤 Base64 block。读取 JSONL 时不会自动还原附件图像；再次处理图像仍需提交原始消息。 +如果模型调用失败，或判断无需写入记忆卡片，已经保存的附件仍然保留，不会自动清理。本地来源路径也可以交给 `read_image` 读取。 ## 消息时间 diff --git a/reme/config/default.yaml b/reme/config/default.yaml index 8da90d01..e162d044 100644 --- a/reme/config/default.yaml +++ b/reme/config/default.yaml @@ -184,7 +184,7 @@ jobs: description: "optional hint" include_images: type: boolean - description: "Use session images with a caller-configured vision-capable AgentScope model and formatter" + description: "Use session images with a vision-capable AgentScope model; save Base64 originals and cite sources" default: false date: type: string diff --git a/reme/steps/evolve/auto_memory.py b/reme/steps/evolve/auto_memory.py index ffb509a5..c1968a9f 100644 --- a/reme/steps/evolve/auto_memory.py +++ b/reme/steps/evolve/auto_memory.py @@ -1,5 +1,7 @@ """auto_memory — record conversation facts into a daily note via an agent.""" +import base64 +import binascii import datetime from pathlib import Path import re @@ -14,16 +16,115 @@ from agentscope.message import DataBlock, Msg, TextBlock, UserMsg from ._evolve import agent_reply_result_text, format_history, now from ..base_step import BaseStep -from ..file_io import extract_daily_date, parse_daily_date, refresh_day_index +from ..file_io import extract_daily_date, get_path_lock, parse_daily_date, refresh_day_index, write_file_safe from ..file_io import validate_filename_component, validate_session_id +from ..file_io._path import IMAGE_MIME_BY_EXT, _check_path_permission, resolve_path from ..index import normalize_posix_path from ...components import R +from ...utils.wikilink_handler import WikilinkHandler _SESSION_ID_KEY = "session_id" _SOURCE_CONVERSATION_KEY = "source_conversation" _MESSAGE_TIME_ALIASES = ("time_created", "timestamp", "createdAt", "timeCreated", "created_time") +def _checked_write_path(workspace: Path, relative: str, allowed_paths) -> Path: + """Use the file Jobs' path and permission checks for attachments and note metadata.""" + target, error = resolve_path(workspace, relative) + if error or target is None: + raise ValueError(error or "Invalid write path") + if not _check_path_permission(workspace, target, allowed_paths): + raise PermissionError(f"No permission to write {relative}") + return target + + +def _check_existing_image(target: Path) -> bool: + if target.is_file(): + return True + if target.exists(): + raise ValueError(f"Image attachment path is not a file: {target.name}") + return False + + +async def _save_session_images( + workspace: Path, + session_dir: str, + session_id: str, + messages: list[Msg], + images: dict[str, DataBlock], + allowed_paths=None, +) -> list[str]: + """Save Base64 bytes and annotate prepared image markers; URLs remain remote. + + ``messages`` must be the invocation-owned copies from image preparation. + Image data and caller-owned messages remain unchanged. Image identity uses + the message ID and block position, not the image block's optional identity. + """ + workspace = workspace.resolve() + error = validate_filename_component(session_id, kind="session_id") + if error: + raise ValueError(error) + if Path(session_dir).is_absolute(): + raise ValueError("session_dir must be workspace-relative") + pending: dict[Path, tuple[str, bytes]] = {} + replacements = [] + sources = [] + for message in messages: + for index, block in enumerate(message.content): + if not isinstance(block, TextBlock) or block.text not in images: + continue + image = images[block.text] + source = image.source + if source.type == "url": + reference = str(source.url) + else: + # A reversible encoding, not a hash: also distinguish IDs on + # case-insensitive filesystems without trusting IDs as paths. + encoded_id = message.id.encode("utf-8").hex() + suffix = next( + (ext for ext, mime in IMAGE_MIME_BY_EXT.items() if mime == source.media_type.lower()), + ".bin", + ) + filename = f"msg-{encoded_id}-image-{index}{suffix}" + if len(filename) > 255: + raise ValueError("Message ID is too long for an image attachment filename") + relative = (Path(session_dir) / "images" / session_id / filename).as_posix() + error = WikilinkHandler.validate_src_dst(relative, relative) + if error: + raise ValueError(f"Image attachment path cannot be linked: {error}") + target = _checked_write_path(workspace, relative, allowed_paths) + # Wait for any in-process writer before reusing the attachment. + # Stable message IDs and block positions identify an image. + async with await get_path_lock(target): + target = _checked_write_path(workspace, relative, allowed_paths) + if not _check_existing_image(target): + try: + payload = base64.b64decode(source.data, validate=True) + except (ValueError, binascii.Error) as exc: + raise ValueError("Image source contains invalid Base64") from exc + if not payload: + raise ValueError("Image source contains empty Base64 data") + if target in pending and pending[target][1] != payload: + raise ValueError("Conflicting images share the same message ID and block position") + pending[target] = (relative, payload) + reference = f"[[{relative}]]" + replacements.append((message, index, block, reference)) + if reference not in sources: + sources.append(reference) + + # Validate every input before creating files. Recheck each path under the + # existing file-operation lock so concurrent saves cannot overwrite it. + for relative, payload in pending.values(): + target = _checked_write_path(workspace, relative, allowed_paths) + async with await get_path_lock(target): + target = _checked_write_path(workspace, relative, allowed_paths) + if not _check_existing_image(target): + await write_file_safe(target, payload) + for message, index, block, reference in replacements: + message.content[index] = block.model_copy(update={"text": f"Image source: {reference}\n{block.text}"}) + return sources + + def _sanitize_msg_for_save(msg: Msg) -> Msg: new_content = [] changed = False @@ -122,11 +223,39 @@ class AutoMemoryStep(BaseStep): notes = list_response.metadata.get("notes") or [] return self._find_session_note(notes, session_id) - async def _ensure_session_frontmatter(self, path: str, session_id: str) -> None: + async def _ensure_session_frontmatter( + self, + path: str, + session_id: str, + image_sources: list[str] | None = None, + ) -> None: metadata = { _SESSION_ID_KEY: session_id, _SOURCE_CONVERSATION_KEY: self._session_link(session_id), } + if image_sources is not None and ( + not isinstance(image_sources, list) or any(not isinstance(source, str) for source in image_sources) + ): + raise ValueError("source_images must be a list of strings") + if image_sources: + workspace = self.file_store.workspace_path.resolve() + allowed_paths = self.context.get("_allowed_paths") + target = _checked_write_path(workspace, path, allowed_paths) + # Read and merge under the same lock as native file Jobs so a + # concurrent write cannot lose its source links to a stale list. + async with await get_path_lock(target): + target = _checked_write_path(workspace, path, allowed_paths) + if not target.is_file(): + raise ValueError(f"Memory note not found: {path}") + post = frontmatter.loads(target.read_text(encoding="utf-8")) + current_sources = post.metadata.get("source_images", []) + if not isinstance(current_sources, list) or any(not isinstance(s, str) for s in current_sources): + raise ValueError("Existing source_images must be a list of strings") + metadata["source_images"] = list(dict.fromkeys([*current_sources, *image_sources])) + if not all(post.metadata.get(key) == value for key, value in metadata.items()): + post.metadata.update(metadata) + await write_file_safe(target, frontmatter.dumps(post)) + return current = self._frontmatter(path) if all(current.get(key) == value for key, value in metadata.items()): return @@ -317,9 +446,9 @@ class AutoMemoryStep(BaseStep): prepared[message_index].content[block_index] = TextBlock(text=marker) return prepared, image_blocks, reply_kwargs - @staticmethod - def _image_user_message(prompt: str, images: dict[str, DataBlock]) -> UserMsg: - """Restore images after the existing templates and history hooks have rendered.""" + def _image_user_message(self, prompt: str, images: dict[str, DataBlock]) -> UserMsg: + """Add source instructions and restore images after rendering the memory prompt.""" + prompt += "\n\n" + self.prompt_format("image_sources_instructions") parts = re.split("(" + "|".join(map(re.escape, images)) + ")", prompt) if [part for part in parts if part in images] != list(images): raise ValueError("Memory prompt must preserve every image once in conversation order") @@ -389,6 +518,24 @@ class AutoMemoryStep(BaseStep): created = note is None before_note_path = note_path before_note_bytes = self._note_bytes(note_path) if note_path else None + # Keep attachment provenance even when the Agent uses a full rewrite, + # or this call updates an image-backed note with images disabled. + image_sources = [] + if before_note_bytes is not None: + previous_sources = frontmatter.loads(before_note_bytes.decode("utf-8")).get("source_images", []) + if isinstance(previous_sources, list) and all(isinstance(source, str) for source in previous_sources): + image_sources = previous_sources + elif images: + raise ValueError("Existing source_images must be a list of strings") + if images: + image_sources += await _save_session_images( + self.file_store.workspace_path, + self._session_dir(), + session_id, + history_messages, + images, + self.context.get("_allowed_paths"), + ) self.logger.info( f"[{self.name}] note lookup session_id={session_id!r} path={note_path!r} " f"created={created} msgs={len(messages)} hint={bool(memory_hint)}", @@ -442,9 +589,11 @@ class AutoMemoryStep(BaseStep): self.logger.info(f"[{self.name}] done without note session_id={session_id!r} modified=False") return note_path = str(note["path"]) + if image_sources: + await self._ensure_session_frontmatter(note_path, session_id, image_sources) else: try: - await self._ensure_session_frontmatter(note_path, session_id) + await self._ensure_session_frontmatter(note_path, session_id, image_sources) note_path = await self._rename_from_frontmatter_name(note_path, day) except RuntimeError as exc: self.context.response.success = False diff --git a/reme/steps/evolve/auto_memory.yaml b/reme/steps/evolve/auto_memory.yaml index 0767a502..3cc7e102 100644 --- a/reme/steps/evolve/auto_memory.yaml +++ b/reme/steps/evolve/auto_memory.yaml @@ -45,6 +45,19 @@ system_prompt_zh: | - `description` = 详细总结；模糊的描述如 "notes" / "misc" 不可接受。 - **永远不要设置 `status`**——它是下游处理保留的字段。 +image_sources_instructions: | + ## Image References + + Each image has an `Image source:` immediately before it. Read it together with the surrounding conversation, keeping the speaker, time and text-image order. + When recording a visual fact, cite that image's exact source in the same sentence or bullet. Use the supplied wikilink for a saved attachment, or a Markdown link for an HTTP(S) source. When a fact compares several images, cite each relevant source. + Keep existing image citations when updating a fact that is still retained. A list of image links at the end, or the `source_images` frontmatter, does not replace these contextual citations. +image_sources_instructions_zh: | + ## 图像引用 + + 每张图片前紧邻的 `Image source:` 是它的来源。结合前后对话理解图片，保留说话人、时间和图文顺序。 + 记录视觉事实时，在对应句子或条目中引用该图的准确来源。保存的附件使用提供的双括号链接，HTTP(S) 来源使用 Markdown 链接；涉及多张图的比较时引用各自的来源。 + 更新时保留仍有效事实对应的图像引用。文末的图片列表或 `source_images` frontmatter 不能替代正文中的对应引用。 + user_message_create: | Today: {today} Extra hint: {note} diff --git a/tests/unit/test_auto_memory_direct_images.py b/tests/unit/test_auto_memory_direct_images.py index d6f9f339..2b63af16 100644 --- a/tests/unit/test_auto_memory_direct_images.py +++ b/tests/unit/test_auto_memory_direct_images.py @@ -5,12 +5,14 @@ import base64 import copy from pathlib import Path +import re from types import SimpleNamespace from unittest.mock import AsyncMock, Mock from agentscope.formatter import DashScopeChatFormatter, OpenAIChatFormatter from agentscope.message import Base64Source, DataBlock, Msg, TextBlock, URLSource import httpx +import frontmatter import pytest import yaml @@ -23,6 +25,7 @@ from reme.components.job import BaseJob from reme.components.tag_index import LocalTagIndex from reme.schema import ApplicationConfig from reme.steps.evolve.auto_memory import AutoMemoryStep +from reme.steps.file_io import WriteStep from .test_auto_tag import _TaggingWrapper, _write_note @@ -31,10 +34,10 @@ _SESSION = "image-input" def _image(source=None): - # Source validation/decoding belongs to the formatter/provider, not this Step. + # Attachment storage decodes Base64; pixel decoding still belongs to the provider. return DataBlock( id="duplicate-image-id", - source=source or Base64Source(media_type="image/png", data="not-decoded-by-ReMe"), + source=source or Base64Source(media_type="image/png", data=base64.b64encode(b"original image bytes").decode()), ) @@ -171,8 +174,7 @@ async def test_sources_interleave_unchanged_through_native_formatters(setup, mon message.content.extend([TextBlock(text="Between images."), _image(URLSource(media_type="image/png", url=url))]) message.content.append(TextBlock(text="After both images.")) original = copy.deepcopy(message.model_dump()) - forbidden = Mock(side_effect=AssertionError("Auto Memory must not download or decode image sources")) - monkeypatch.setattr(base64, "b64decode", forbidden) + forbidden = Mock(side_effect=AssertionError("Auto Memory must not download image sources")) monkeypatch.setattr(httpx.AsyncClient, "send", forbidden) monkeypatch.setattr(httpx.Client, "send", forbidden) @@ -185,7 +187,11 @@ async def test_sources_interleave_unchanged_through_native_formatters(setup, mon assert "Remember this observation." in inputs.content[0].text assert inputs.get_text_content().count("[Image 1]") == 1 assert "__reme_image_" not in inputs.get_text_content() - assert inputs.content[2].text.strip() == "Between images." + assert inputs.content[2].text.lstrip().startswith("Between images.") + assert "Image source: [[session/images/image-input/" in inputs.content[0].text + assert f"Image source: {url}" in inputs.content[2].text + assets = list((path.parents[1] / "images" / _SESSION).glob("*.png")) + assert len(assets) == 1 and assets[0].read_bytes() == b"original image bytes" assert inputs.content[-1].text.index("After both images.") < inputs.content[-1].text.index("# Your Task") assert [inputs.content[index].model_dump() for index in (1, 3)] == [ message.content[index].model_dump() for index in (1, 3) @@ -292,6 +298,7 @@ async def test_full_history_hook_image_only_turn_and_update_boundaries(setup, mo step.prompt.language = language note_path = f"daily/{_DAY}/existing.md" if existing: + _write_note(step.file_store.workspace_path / note_path) monkeypatch.setattr(step, "_list_session_note", AsyncMock(return_value={"path": note_path})) monkeypatch.setattr(step, "_ensure_session_frontmatter", AsyncMock()) monkeypatch.setattr(step, "_rename_from_frontmatter_name", AsyncMock(return_value=note_path)) @@ -452,3 +459,134 @@ async def test_default_job_still_passes_memory_changes_to_auto_tag(setup, monkey assert "auto_memory_images" not in response.metadata assert isinstance(wrapper.reply.call_args.args[0], Msg) assert len(tagger.calls) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("enabled_update", [False, True]) +async def test_image_links_survive_native_full_rewrite_and_keep_context(setup, monkeypatch, enabled_update): + step, wrapper, session_path = setup + workspace = step.file_store.workspace_path + note_path = f"daily/{_DAY}/image-memory.md" + target = workspace / note_path + old_link = "[[user-owned/manual.png]]" + writer = WriteStep(app_context=step.app_context, file_store=step.file_store) + await writer(path=note_path, name="image-memory", content="Existing fact.", metadata={"source_images": [old_link]}) + + monkeypatch.setattr(step, "_list_session_note", AsyncMock(return_value={"path": note_path})) + monkeypatch.setattr("reme.steps.evolve.auto_memory.refresh_day_index", AsyncMock(return_value={})) + first, second = _message(), _message("second", timestamp=f"{_DAY}T11:00:00") + first.content[0].text = "Alice's picture follows." + second.name = "Bob" + second.content[0].text = "Bob's picture follows." + second.content[1].source.data = base64.b64encode(b"Bob's different image").decode() + before = [message.model_dump() for message in (first, second)] + written_bodies = [] + + async def full_rewrite(inputs, **kwargs): + assert kwargs["injected_job_kwargs"] == {"_allowed_paths": [note_path]} + if isinstance(inputs, Msg): + blocks = inputs.content + assert [block.type for block in blocks] == ["text", "data", "text", "data", "text"] + citations = [] + for text_index, owner, image in ((0, "Alice", first.content[1]), (2, "Bob", second.content[1])): + assert f"{owner}'s picture follows." in blocks[text_index].text + assert blocks[text_index + 1].model_dump() == image.model_dump() + link = re.search(r"Image source: (\[\[.*?\]\])", blocks[text_index].text).group(1) + assert (workspace / link[2:-2]).read_bytes() == base64.b64decode(image.source.data) + citations.append(f"{owner}'s visual fact: {link}") + assert "same sentence or bullet" in blocks[-1].text + body = "\n".join(citations) + else: + body = written_bodies[0] + "\nUpdated text-only fact." + written_bodies.append(body) + # Native write replaces frontmatter as well as the body. Auto Memory + # must restore previous image sources without rewriting this body. + await writer(path=note_path, name="image-memory", description="Images", content=body) + return {"result": "Updated."} + + wrapper.reply.side_effect = full_rewrite + await _run(step, [first, second], include_images=True) + initial = frontmatter.loads(target.read_text()) + sources = initial["source_images"] + assert sources[0] == old_link and len(sources) == 3 + assert initial.content == written_bodies[0] + attachments = {path: path.stat().st_mtime_ns for path in (workspace / "session/images").rglob("*.png")} + assert len(attachments) == 2 + + await _run(step, [first, second], include_images=enabled_update) + + updated = frontmatter.loads(target.read_text()) + assert updated["source_images"] == sources + assert updated.content == written_bodies[-1] + assert {path: path.stat().st_mtime_ns for path in attachments} == attachments + assert session_path.read_bytes() == _saved_line(first) + _saved_line(second) + assert [message.model_dump() for message in (first, second)] == before + assert step.context["changes"] == ([] if enabled_update else [{"change": "modified", "path": note_path}]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("value", ["manual", None, [7]]) +@pytest.mark.parametrize("enabled,images", [(True, True), (False, True), (True, False)]) +async def test_invalid_note_sources_only_block_image_writes(setup, monkeypatch, value, enabled, images): + step, wrapper, session_path = setup + note_path = f"daily/{_DAY}/existing.md" + target = step.file_store.workspace_path / note_path + _write_note(target) + post = frontmatter.loads(target.read_text()) + post["source_images"] = value + target.write_text(frontmatter.dumps(post)) + before = target.read_bytes() + monkeypatch.setattr(step, "_list_session_note", AsyncMock(return_value={"path": note_path})) + monkeypatch.setattr(step, "_ensure_session_frontmatter", AsyncMock()) + monkeypatch.setattr(step, "_rename_from_frontmatter_name", AsyncMock(return_value=note_path)) + monkeypatch.setattr("reme.steps.evolve.auto_memory.refresh_day_index", AsyncMock(return_value={})) + message = _message(images=images) + if enabled and images: + with pytest.raises(ValueError, match="source_images"): + await _run(step, [message], include_images=enabled) + wrapper.reply.assert_not_called() + else: + response = await _run(step, [message], include_images=enabled) + assert response.success + assert isinstance(wrapper.reply.call_args.args[0], str) + assert session_path.read_bytes() == _saved_line(message) + assert not (step.file_store.workspace_path / "session/images").exists() + assert target.read_bytes() == before + + +@pytest.mark.asyncio +@pytest.mark.parametrize("enabled", [False, True]) +@pytest.mark.parametrize("lookup_fails", [False, True]) +async def test_source_save_precedes_single_note_lookup_and_image_writes(setup, monkeypatch, enabled, lookup_fails): + step, wrapper, session_path = setup + note_path = f"daily/{_DAY}/existing.md" + _write_note(step.file_store.workspace_path / note_path) + message = _message() + image_dir = step.file_store.workspace_path / "session/images" + + async def find_note(*_args): + assert session_path.read_bytes() == _saved_line(message) + assert not image_dir.exists() + wrapper.reply.assert_not_called() + if lookup_fails: + raise RuntimeError("daily_list failed: unavailable") + return {"path": note_path} + + lookup = AsyncMock(side_effect=find_note) + monkeypatch.setattr(step, "_list_session_note", lookup) + monkeypatch.setattr(step, "_ensure_session_frontmatter", AsyncMock()) + monkeypatch.setattr(step, "_rename_from_frontmatter_name", AsyncMock(return_value=note_path)) + monkeypatch.setattr("reme.steps.evolve.auto_memory.refresh_day_index", AsyncMock(return_value={})) + response = await _run(step, [message], include_images=enabled) + + lookup.assert_awaited_once_with(_DAY, _SESSION) + assert response.success is not lookup_fails + assert session_path.read_bytes() == _saved_line(message) + assert image_dir.exists() is (enabled and not lookup_fails) + if lookup_fails: + assert response.answer == "daily_list failed: unavailable" + assert response.metadata == {"date": _DAY, "modified": False, "n_messages": 1} + wrapper.reply.assert_not_called() + else: + wrapper.reply.assert_awaited_once() + assert isinstance(wrapper.reply.call_args.args[0], Msg if enabled else str) diff --git a/tests/unit/test_session_image_assets.py b/tests/unit/test_session_image_assets.py new file mode 100644 index 00000000..5d03ab11 --- /dev/null +++ b/tests/unit/test_session_image_assets.py @@ -0,0 +1,236 @@ +"""Session attachments preserve bytes, identity, file boundaries and user links.""" + +# pylint: disable=protected-access,missing-function-docstring + +import asyncio +import base64 +import importlib +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock + +from agentscope.message import Base64Source, DataBlock, Msg, TextBlock, URLSource +import frontmatter +import httpx +import pytest + +from reme.components.runtime_context import RuntimeContext +from reme.steps.evolve.auto_memory import AutoMemoryStep +from reme.steps.evolve.auto_memory import _save_session_images as save_session_images +from reme.steps.file_io.write import WriteStep + + +def _prepared(*sources, message_id="message-a"): + content = [TextBlock(text="before")] + for source in sources: + if not isinstance(source, URLSource): + source = Base64Source(media_type="image/png", data=base64.b64encode(source).decode()) + content.append(DataBlock(id="same-block-id", source=source)) + content.append(TextBlock(text="after")) + original = Msg(id=message_id, name="Alice", role="user", content=content) + prepared = original.model_copy(deep=True) + images = {} + for index, block in enumerate(prepared.content): + if isinstance(block, DataBlock): + marker = f"__image_{index}__" + images[marker] = block + prepared.content[index] = TextBlock(text=marker) + return original, [prepared], images + + +@pytest.fixture(name="memory_step") +def session_memory_step(tmp_path): + step = AutoMemoryStep() + step.file_store = SimpleNamespace(workspace_path=tmp_path) + step.context = RuntimeContext() + return step + + +@pytest.mark.asyncio +async def test_bytes_positions_and_existing_images_are_reused_without_reading(tmp_path, monkeypatch): + original, messages, images = _prepared(b"first image", b"second image") + snapshot = original.model_dump() + sources = await save_session_images(tmp_path, "sessions", "chat", messages, images) + assert len(sources) == 2 + paths = [tmp_path / source[2:-2] for source in sources] + assert [path.read_bytes() for path in paths] == [b"first image", b"second image"] + assert [path.name.rsplit("-", 1)[1] for path in paths] == ["1.png", "2.png"] + assert messages[0].content[1].text == f"Image source: {sources[0]}\n__image_1__" + assert messages[0].content[2].text == f"Image source: {sources[1]}\n__image_2__" + assert original.model_dump() == snapshot + assert [block.model_dump() for block in images.values()] == snapshot["content"][1:3] + mtimes = [path.stat().st_mtime_ns for path in paths] + _, replay, replay_images = _prepared(b"different first image", b"different second image") + forbidden = Mock(side_effect=AssertionError("Existing attachments must not be decoded or read")) + with monkeypatch.context() as patch: + patch.setattr(base64, "b64decode", forbidden) + patch.setattr(Path, "read_bytes", forbidden) + assert await save_session_images(tmp_path, "sessions", "chat", replay, replay_images) == sources + forbidden.assert_not_called() + assert [path.read_bytes() for path in paths] == [b"first image", b"second image"] + assert [path.stat().st_mtime_ns for path in paths] == mtimes + + _, extended, extended_images = _prepared(b"different first image", b"different second image", b"third image") + decoder = Mock(wraps=base64.b64decode) + with monkeypatch.context() as patch: + patch.setattr(base64, "b64decode", decoder) + patch.setattr(Path, "read_bytes", forbidden) + extended_sources = await save_session_images(tmp_path, "sessions", "chat", extended, extended_images) + decoder.assert_called_once_with(extended_images["__image_3__"].source.data, validate=True) + assert extended_sources[:2] == sources and len(extended_sources) == 3 + assert (tmp_path / extended_sources[2][2:-2]).read_bytes() == b"third image" + assert [path.stat().st_mtime_ns for path in paths] == mtimes + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "media_type,suffix", + [ + ("image/png", ".png"), + ("image/jpeg", ".jpg"), + ("image/webp", ".webp"), + ("image/gif", ".gif"), + ("image/bmp", ".bmp"), + ("image/tiff", ".tiff"), + ("image/heic", ".heic"), + ("image/avif", ".bin"), + ], +) +async def test_image_suffixes_preserve_bytes_including_unknown_types(tmp_path, media_type, suffix): + _, messages, images = _prepared(b"unchanged bytes") + images["__image_1__"].source.media_type = media_type + sources = await save_session_images(tmp_path, "session", "chat", messages, images) + path = tmp_path / sources[0][2:-2] + assert path.suffix == suffix and path.read_bytes() == b"unchanged bytes" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("message_id", ["../../outside", "A", "a", "图像/a%"]) +async def test_message_id_is_encoded_not_used_as_a_path(tmp_path, message_id): + _, messages, images = _prepared(b"image", message_id=message_id) + sources = await save_session_images(tmp_path, "session", "chat", messages, images) + path = tmp_path / sources[0][2:-2] + encoded = path.name.removeprefix("msg-").split("-image-", 1)[0] + assert bytes.fromhex(encoded).decode() == message_id + assert path.parent == tmp_path / "session/images/chat" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "session_dir,session_id", + [("session", "chat#1"), ("session", "[chat]"), ("sessions#1", "chat"), ("[sessions]", "chat")], +) +async def test_unlinkable_attachment_paths_fail_before_writing(tmp_path, session_dir, session_id): + _, messages, images = _prepared(b"image") + with pytest.raises(ValueError, match="cannot be linked"): + await save_session_images(tmp_path, session_dir, session_id, messages, images) + assert not list(tmp_path.iterdir()) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", ["base64", "denied", "directory", "symlink", "traversal"]) +async def test_invalid_inputs_are_checked_before_any_image_write(tmp_path, failure): + _, messages, images = _prepared(b"first", b"second") + session_dir = "session" + allowed_paths = None + second = tmp_path / f"session/images/chat/msg-{'message-a'.encode().hex()}-image-2.png" + if failure == "base64": + images["__image_2__"].source.data = "invalid base64!" + elif failure == "denied": + allowed_paths = ["unrelated"] + elif failure == "directory": + second.mkdir(parents=True) + elif failure == "symlink": + (tmp_path / "session").symlink_to(tmp_path.parent, target_is_directory=True) + else: + session_dir = "../outside" + with pytest.raises((ValueError, PermissionError)): + await save_session_images(tmp_path, session_dir, "chat", messages, images, allowed_paths) + assert messages[0].content[1].text == "__image_1__" + assert not list(tmp_path.glob("session/images/chat/*-image-1.png")) + + +@pytest.mark.asyncio +async def test_http_sources_remain_urls_without_local_io(tmp_path, monkeypatch): + def fail_network(*_args, **_kwargs): + pytest.fail("Auto Memory must not download image URLs") + + monkeypatch.setattr(httpx.Client, "request", fail_network) + monkeypatch.setattr(httpx.AsyncClient, "request", fail_network) + url = "https://example.com/image.png?token=original" + _, messages, images = _prepared(URLSource(url=url, media_type="image/png")) + assert await save_session_images(tmp_path, "session", "chat", messages, images) == [url] + assert messages[0].content[1].text == f"Image source: {url}\n__image_1__" + assert not list(tmp_path.iterdir()) + + +@pytest.mark.asyncio +async def test_source_links_preserve_body_unknown_links_and_noop_mtime(tmp_path, memory_step): + path = tmp_path / "note.md" + path.write_text("---\nsource_images: ['[[manual.md]]']\nuser_owned: keep\n---\n\n# Body\n\nKeep this text.") + before = frontmatter.loads(path.read_text()) + await memory_step._ensure_session_frontmatter("note.md", "chat", ["[[new.png]]", "[[new.png]]", "[[manual.md]]"]) + after = frontmatter.loads(path.read_text()) + assert after.content == before.content + assert after.metadata == { + "source_images": ["[[manual.md]]", "[[new.png]]"], + "user_owned": "keep", + "session_id": "chat", + "source_conversation": "[[session/dialog/chat.jsonl]]", + } + mtime = path.stat().st_mtime_ns + await memory_step._ensure_session_frontmatter("note.md", "chat", ["[[new.png]]"]) + assert path.stat().st_mtime_ns == mtime + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "existing,value", + [(True, "manual"), (True, None), (True, [3]), (False, "[[image.png]]"), (False, [None])], +) +async def test_invalid_sources_do_not_rewrite_note(tmp_path, memory_step, existing, value): + path = tmp_path / "note.md" + path.write_text(frontmatter.dumps(frontmatter.Post("User body.", **({"source_images": value} if existing else {})))) + before = path.read_bytes() + with pytest.raises(ValueError, match="list of strings"): + await memory_step._ensure_session_frontmatter("note.md", "chat", ["[[new.png]]"] if existing else value) + assert path.read_bytes() == before + + +@pytest.mark.asyncio +async def test_source_frontmatter_rereads_after_concurrent_native_write(tmp_path, monkeypatch, memory_step): + monkeypatch.chdir(tmp_path) + write_module = importlib.import_module("reme.steps.file_io.write") + native_write = write_module.write_file_safe + entered, release = asyncio.Event(), asyncio.Event() + + async def delayed_write(*args, **kwargs): + entered.set() + await release.wait() + await native_write(*args, **kwargs) + + monkeypatch.setattr(write_module, "write_file_safe", delayed_write) + path = tmp_path / "note.md" + path.write_text("Old content") + writer = WriteStep(file_store=SimpleNamespace(workspace_path=tmp_path)) + context = RuntimeContext( + path="note.md", + content="New user content", + metadata={"source_images": ["[[manual.md]]"], "user_owned": 7}, + ) + writing = asyncio.create_task(writer(context)) + await asyncio.wait_for(entered.wait(), timeout=5) + merging = asyncio.create_task(memory_step._ensure_session_frontmatter("note.md", "chat", ["[[image.png]]"])) + await asyncio.sleep(0) + assert not merging.done() + release.set() + await asyncio.wait_for(asyncio.gather(writing, merging), timeout=5) + assert context.response.success + post = frontmatter.loads(path.read_text()) + assert post.content == "New user content" + assert post.metadata == { + "source_images": ["[[manual.md]]", "[[image.png]]"], + "user_owned": 7, + "session_id": "chat", + "source_conversation": "[[session/dialog/chat.jsonl]]", + }