diff --git a/adapter/matrix/store.py b/adapter/matrix/store.py index 30ee076..d046640 100644 --- a/adapter/matrix/store.py +++ b/adapter/matrix/store.py @@ -7,6 +7,8 @@ USER_META_PREFIX = "matrix_user:" ROOM_STATE_PREFIX = "matrix_state:" SKILLS_MSG_PREFIX = "matrix_skills_msg:" PENDING_CONFIRM_PREFIX = "matrix_pending_confirm:" +LOAD_PENDING_PREFIX = "matrix_load_pending:" +RESET_PENDING_PREFIX = "matrix_reset_pending:" async def get_room_meta(store: StateStore, room_id: str) -> dict | None: @@ -74,3 +76,40 @@ async def clear_pending_confirm( store: StateStore, user_id: str, room_id: str | None = None ) -> None: await store.delete(_pending_confirm_key(user_id, room_id)) + + +def _load_pending_key(user_id: str, room_id: str) -> str: + return f"{LOAD_PENDING_PREFIX}{user_id}:{room_id}" + + +async def get_load_pending(store: StateStore, user_id: str, room_id: str) -> dict | None: + return await store.get(_load_pending_key(user_id, room_id)) + + +async def set_load_pending(store: StateStore, user_id: str, room_id: str, data: dict) -> None: + await store.set(_load_pending_key(user_id, room_id), data) + + +async def clear_load_pending(store: StateStore, user_id: str, room_id: str) -> None: + await store.delete(_load_pending_key(user_id, room_id)) + + +def _reset_pending_key(user_id: str, room_id: str) -> str: + return f"{RESET_PENDING_PREFIX}{user_id}:{room_id}" + + +async def get_reset_pending(store: StateStore, user_id: str, room_id: str) -> dict | None: + return await store.get(_reset_pending_key(user_id, room_id)) + + +async def set_reset_pending( + store: StateStore, + user_id: str, + room_id: str, + data: dict, +) -> None: + await store.set(_reset_pending_key(user_id, room_id), data) + + +async def clear_reset_pending(store: StateStore, user_id: str, room_id: str) -> None: + await store.delete(_reset_pending_key(user_id, room_id)) diff --git a/sdk/prototype_state.py b/sdk/prototype_state.py index ccb75f1..6374982 100644 --- a/sdk/prototype_state.py +++ b/sdk/prototype_state.py @@ -31,6 +31,8 @@ class PrototypeStateStore: def __init__(self) -> None: self._users: dict[str, User] = {} self._settings: dict[str, dict[str, Any]] = {} + self._saved_sessions: dict[str, list[dict[str, str]]] = {} + self._last_tokens_used: dict[str, int] = {} async def get_or_create_user( self, @@ -78,3 +80,16 @@ class PrototypeStateStore: elif action.action == "set_safety": safety = settings.setdefault("safety", DEFAULT_SAFETY.copy()) safety[action.payload["trigger"]] = action.payload.get("enabled", True) + + async def add_saved_session(self, user_id: str, name: str) -> None: + sessions = self._saved_sessions.setdefault(user_id, []) + sessions.append({"name": name, "created_at": datetime.now(UTC).isoformat()}) + + async def list_saved_sessions(self, user_id: str) -> list[dict[str, str]]: + return list(self._saved_sessions.get(user_id, [])) + + async def get_last_tokens_used(self, user_id: str) -> int: + return self._last_tokens_used.get(user_id, 0) + + async def set_last_tokens_used(self, user_id: str, tokens: int) -> None: + self._last_tokens_used[user_id] = tokens diff --git a/tests/platform/test_prototype_state.py b/tests/platform/test_prototype_state.py index b5f5dc3..e42f650 100644 --- a/tests/platform/test_prototype_state.py +++ b/tests/platform/test_prototype_state.py @@ -89,3 +89,46 @@ async def test_update_settings_supports_toggle_skill_and_setters(): assert settings.skills["web-search"] is True assert settings.soul["instructions"] == "Be concise" assert settings.safety["social-post"] is False + + +@pytest.mark.asyncio +async def test_add_saved_session_appends_named_entries(): + store = PrototypeStateStore() + + await store.add_saved_session("usr-matrix-@alice:example.org", "alpha") + await store.add_saved_session("usr-matrix-@alice:example.org", "beta") + + sessions = await store.list_saved_sessions("usr-matrix-@alice:example.org") + + assert [session["name"] for session in sessions] == ["alpha", "beta"] + assert all("created_at" in session for session in sessions) + + +@pytest.mark.asyncio +async def test_list_saved_sessions_returns_copy(): + store = PrototypeStateStore() + + await store.add_saved_session("usr-matrix-@alice:example.org", "alpha") + + sessions = await store.list_saved_sessions("usr-matrix-@alice:example.org") + sessions.append({"name": "tampered", "created_at": "never"}) + + stored = await store.list_saved_sessions("usr-matrix-@alice:example.org") + + assert [session["name"] for session in stored] == ["alpha"] + + +@pytest.mark.asyncio +async def test_get_last_tokens_used_defaults_to_zero(): + store = PrototypeStateStore() + + assert await store.get_last_tokens_used("usr-matrix-@alice:example.org") == 0 + + +@pytest.mark.asyncio +async def test_set_last_tokens_used_persists_value(): + store = PrototypeStateStore() + + await store.set_last_tokens_used("usr-matrix-@alice:example.org", 321) + + assert await store.get_last_tokens_used("usr-matrix-@alice:example.org") == 321