feat: add matrix staged attachment state
This commit is contained in:
parent
105ecc68ed
commit
0eaf124e21
2 changed files with 188 additions and 0 deletions
|
|
@ -1,5 +1,8 @@
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
from weakref import WeakValueDictionary
|
||||||
|
|
||||||
from core.store import StateStore
|
from core.store import StateStore
|
||||||
|
|
||||||
ROOM_META_PREFIX = "matrix_room:"
|
ROOM_META_PREFIX = "matrix_room:"
|
||||||
|
|
@ -9,6 +12,8 @@ SKILLS_MSG_PREFIX = "matrix_skills_msg:"
|
||||||
PENDING_CONFIRM_PREFIX = "matrix_pending_confirm:"
|
PENDING_CONFIRM_PREFIX = "matrix_pending_confirm:"
|
||||||
LOAD_PENDING_PREFIX = "matrix_load_pending:"
|
LOAD_PENDING_PREFIX = "matrix_load_pending:"
|
||||||
RESET_PENDING_PREFIX = "matrix_reset_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:
|
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:
|
async def clear_reset_pending(store: StateStore, user_id: str, room_id: str) -> None:
|
||||||
await store.delete(_reset_pending_key(user_id, room_id))
|
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))
|
||||||
|
|
|
||||||
|
|
@ -3,14 +3,19 @@ from __future__ import annotations
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from adapter.matrix.store import (
|
from adapter.matrix.store import (
|
||||||
|
STAGED_ATTACHMENTS_PREFIX,
|
||||||
|
add_staged_attachment,
|
||||||
clear_pending_confirm,
|
clear_pending_confirm,
|
||||||
|
clear_staged_attachments,
|
||||||
get_pending_confirm,
|
get_pending_confirm,
|
||||||
get_platform_chat_id,
|
get_platform_chat_id,
|
||||||
get_room_meta,
|
get_room_meta,
|
||||||
get_room_state,
|
get_room_state,
|
||||||
get_skills_message_id,
|
get_skills_message_id,
|
||||||
|
get_staged_attachments,
|
||||||
get_user_meta,
|
get_user_meta,
|
||||||
next_chat_id,
|
next_chat_id,
|
||||||
|
remove_staged_attachment_at,
|
||||||
set_pending_confirm,
|
set_pending_confirm,
|
||||||
set_platform_chat_id,
|
set_platform_chat_id,
|
||||||
set_room_meta,
|
set_room_meta,
|
||||||
|
|
@ -116,3 +121,118 @@ async def test_pending_confirm_roundtrip(store: InMemoryStore):
|
||||||
|
|
||||||
await clear_pending_confirm(store, "!room:m.org")
|
await clear_pending_confirm(store, "!room:m.org")
|
||||||
assert await get_pending_confirm(store, "!room:m.org") is None
|
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) == []
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue