From 0eaf124e21a3ad5a7b80d3931373f1e0e7d77573 Mon Sep 17 00:00:00 2001 From: Mikhail Putilovskij Date: Mon, 20 Apr 2026 16:21:00 +0300 Subject: [PATCH] feat: add matrix staged attachment state --- adapter/matrix/store.py | 68 ++++++++++++++++ tests/adapter/matrix/test_store.py | 120 +++++++++++++++++++++++++++++ 2 files changed, 188 insertions(+) diff --git a/adapter/matrix/store.py b/adapter/matrix/store.py index 5ebb61a..acafa9f 100644 --- a/adapter/matrix/store.py +++ b/adapter/matrix/store.py @@ -1,5 +1,8 @@ from __future__ import annotations +import asyncio +from weakref import WeakValueDictionary + from core.store import StateStore ROOM_META_PREFIX = "matrix_room:" @@ -9,6 +12,8 @@ SKILLS_MSG_PREFIX = "matrix_skills_msg:" PENDING_CONFIRM_PREFIX = "matrix_pending_confirm:" LOAD_PENDING_PREFIX = "matrix_load_pending:" RESET_PENDING_PREFIX = "matrix_reset_pending:" +STAGED_ATTACHMENTS_PREFIX = "matrix_staged_attachments:" +_STAGED_ATTACHMENTS_LOCKS: WeakValueDictionary[str, asyncio.Lock] = WeakValueDictionary() async def get_room_meta(store: StateStore, room_id: str) -> dict | None: @@ -126,3 +131,66 @@ async def set_reset_pending( async def clear_reset_pending(store: StateStore, user_id: str, room_id: str) -> None: await store.delete(_reset_pending_key(user_id, room_id)) + + +def _staged_attachments_key(room_id: str, user_id: str) -> str: + return f"{STAGED_ATTACHMENTS_PREFIX}{room_id}:{user_id}" + + +def _staged_attachments_lock(room_id: str, user_id: str) -> asyncio.Lock: + key = _staged_attachments_key(room_id, user_id) + lock = _STAGED_ATTACHMENTS_LOCKS.get(key) + if lock is None: + lock = asyncio.Lock() + _STAGED_ATTACHMENTS_LOCKS[key] = lock + return lock + + +async def get_staged_attachments( + store: StateStore, room_id: str, user_id: str +) -> list[dict]: + data = await store.get(_staged_attachments_key(room_id, user_id)) + if not isinstance(data, dict): + return [] + + attachments = data.get("attachments") + if not isinstance(attachments, list): + return [] + + return [attachment for attachment in attachments if isinstance(attachment, dict)] + + +async def add_staged_attachment( + store: StateStore, room_id: str, user_id: str, attachment: dict +) -> None: + async with _staged_attachments_lock(room_id, user_id): + attachments = await get_staged_attachments(store, room_id, user_id) + attachments.append(attachment) + await store.set( + _staged_attachments_key(room_id, user_id), {"attachments": attachments} + ) + + +async def remove_staged_attachment_at( + store: StateStore, room_id: str, user_id: str, index: int +) -> dict | None: + async with _staged_attachments_lock(room_id, user_id): + attachments = await get_staged_attachments(store, room_id, user_id) + if index < 0 or index >= len(attachments): + return None + + removed = attachments.pop(index) + if attachments: + await store.set( + _staged_attachments_key(room_id, user_id), {"attachments": attachments} + ) + else: + await store.delete(_staged_attachments_key(room_id, user_id)) + return removed + + +async def clear_staged_attachments( + store: StateStore, room_id: str, user_id: str +) -> None: + async with _staged_attachments_lock(room_id, user_id): + await store.delete(_staged_attachments_key(room_id, user_id)) diff --git a/tests/adapter/matrix/test_store.py b/tests/adapter/matrix/test_store.py index 9fcd2a2..dfb0379 100644 --- a/tests/adapter/matrix/test_store.py +++ b/tests/adapter/matrix/test_store.py @@ -3,14 +3,19 @@ from __future__ import annotations import pytest from adapter.matrix.store import ( + STAGED_ATTACHMENTS_PREFIX, + add_staged_attachment, clear_pending_confirm, + clear_staged_attachments, get_pending_confirm, get_platform_chat_id, get_room_meta, get_room_state, get_skills_message_id, + get_staged_attachments, get_user_meta, next_chat_id, + remove_staged_attachment_at, set_pending_confirm, set_platform_chat_id, set_room_meta, @@ -116,3 +121,118 @@ async def test_pending_confirm_roundtrip(store: InMemoryStore): await clear_pending_confirm(store, "!room:m.org") assert await get_pending_confirm(store, "!room:m.org") is None + + +async def test_staged_attachments_roundtrip(store: InMemoryStore): + room_id = "!room:m.org" + user_id = "@alice:m.org" + + assert await get_staged_attachments(store, room_id, user_id) == [] + + first = {"id": "att-1", "name": "screenshot.png"} + second = {"id": "att-2", "name": "invoice.pdf"} + + await add_staged_attachment(store, room_id, user_id, first) + await add_staged_attachment(store, room_id, user_id, second) + + assert await get_staged_attachments(store, room_id, user_id) == [ + first, + second, + ] + + +@pytest.mark.parametrize( + "stored_value", + [ + None, + "not-a-dict", + [], + 123, + ], +) +async def test_staged_attachments_invalid_container_state_returns_empty_list( + store: InMemoryStore, stored_value, +): + room_id = "!room:m.org" + user_id = "@alice:m.org" + + await store.set(f"{STAGED_ATTACHMENTS_PREFIX}{room_id}:{user_id}", stored_value) + + assert await get_staged_attachments(store, room_id, user_id) == [] + + +async def test_staged_attachments_filters_invalid_entries(store: InMemoryStore): + room_id = "!room:m.org" + user_id = "@alice:m.org" + valid_one = {"id": "att-1", "name": "alpha.png"} + valid_two = {"id": "att-2", "name": "beta.pdf"} + + await store.set( + f"{STAGED_ATTACHMENTS_PREFIX}{room_id}:{user_id}", + { + "attachments": [ + valid_one, + "bad-entry", + None, + {"id": "ignored"}, + valid_two, + ] + }, + ) + + assert await get_staged_attachments(store, room_id, user_id) == [ + valid_one, + {"id": "ignored"}, + valid_two, + ] + + +async def test_staged_attachments_are_scoped_by_room_and_user(store: InMemoryStore): + room_a = "!room-a:m.org" + room_b = "!room-b:m.org" + user_a = "@alice:m.org" + user_b = "@bob:m.org" + + attachment_a = {"id": "att-a", "name": "alpha.png"} + attachment_b = {"id": "att-b", "name": "beta.png"} + attachment_c = {"id": "att-c", "name": "gamma.png"} + + await add_staged_attachment(store, room_a, user_a, attachment_a) + await add_staged_attachment(store, room_a, user_b, attachment_b) + await add_staged_attachment(store, room_b, user_a, attachment_c) + + assert await get_staged_attachments(store, room_a, user_a) == [attachment_a] + assert await get_staged_attachments(store, room_a, user_b) == [attachment_b] + assert await get_staged_attachments(store, room_b, user_a) == [attachment_c] + assert await get_staged_attachments(store, room_b, user_b) == [] + + +async def test_remove_staged_attachment_at_by_zero_based_index( + store: InMemoryStore, +): + room_id = "!room:m.org" + user_id = "@alice:m.org" + first = {"id": "att-1", "name": "first.png"} + second = {"id": "att-2", "name": "second.png"} + third = {"id": "att-3", "name": "third.png"} + + await add_staged_attachment(store, room_id, user_id, first) + await add_staged_attachment(store, room_id, user_id, second) + await add_staged_attachment(store, room_id, user_id, third) + + assert await remove_staged_attachment_at(store, room_id, user_id, 1) == second + assert await get_staged_attachments(store, room_id, user_id) == [first, third] + assert await remove_staged_attachment_at(store, room_id, user_id, 99) is None + assert await remove_staged_attachment_at(store, room_id, user_id, -1) is None + + +async def test_clear_staged_attachments(store: InMemoryStore): + room_id = "!room:m.org" + user_id = "@alice:m.org" + + await add_staged_attachment(store, room_id, user_id, {"id": "att-1"}) + await add_staged_attachment(store, room_id, user_id, {"id": "att-2"}) + + await clear_staged_attachments(store, room_id, user_id) + + assert await get_staged_attachments(store, room_id, user_id) == []