diff --git a/nanobot/session/manager.py b/nanobot/session/manager.py index 4322f799..e787eb6f 100644 --- a/nanobot/session/manager.py +++ b/nanobot/session/manager.py @@ -401,17 +401,17 @@ class SessionManager: @staticmethod def safe_key(key: str) -> str: - """Collision-resistant encoding of a session key for use as a filename stem. + """Public helper used by HTTP handlers to map an arbitrary key to a stable filename stem.""" + return safe_filename(key.replace(":", "_")) - Uses base64url (no padding) so distinct keys always map to distinct - filenames, unlike the previous replace(":", "_") approach which - could collide (e.g. telegram:a_b vs telegram:a:b). - """ + @staticmethod + def _storage_key(key: str) -> str: + """Collision-resistant encoding for internal session storage filenames.""" return base64.urlsafe_b64encode(key.encode()).decode().rstrip("=") def _get_session_path(self, key: str) -> Path: """Get the collision-resistant workspace path for a session.""" - return self.sessions_dir / f"{self.safe_key(key)}.jsonl" + return self.sessions_dir / f"{self._storage_key(key)}.jsonl" def _get_legacy_lossy_path(self, key: str) -> Path: """Previous workspace session path using lossy ':' to '_' replacement.""" diff --git a/nanobot/webui/session_list_index.py b/nanobot/webui/session_list_index.py index 8269500c..7788aef8 100644 --- a/nanobot/webui/session_list_index.py +++ b/nanobot/webui/session_list_index.py @@ -231,8 +231,22 @@ def _indexed_row_for_session(session: Session, path: Path) -> dict[str, Any]: } +def _try_decode_storage_stem(stem: str) -> str | None: + """Try to decode a base64url (no-padding) session storage stem back to the original key.""" + import base64 + + try: + padding = 4 - len(stem) % 4 + if padding != 4: + stem += "=" * padding + return base64.urlsafe_b64decode(stem).decode("utf-8") + except Exception: + return None + + def _scan_session_row(session_manager: SessionManager, path: Path) -> dict[str, Any] | None: - fallback_key = path.stem.replace("_", ":", 1) + storage_key = _try_decode_storage_stem(path.stem) + fallback_key = storage_key or path.stem.replace("_", ":", 1) try: with open(path, encoding="utf-8") as f: first_line = f.readline().strip() diff --git a/tests/agent/test_session_collision.py b/tests/agent/test_session_collision.py index 7484f09a..74d355e2 100644 --- a/tests/agent/test_session_collision.py +++ b/tests/agent/test_session_collision.py @@ -40,8 +40,8 @@ def test_distinct_keys_have_distinct_filenames(tmp_path: Path, monkeypatch) -> N second = sm._get_session_path("telegram:a:b") assert first.name != second.name - assert first.name == f"{SessionManager.safe_key('telegram:a_b')}.jsonl" - assert second.name == f"{SessionManager.safe_key('telegram:a:b')}.jsonl" + assert sm.safe_key("telegram:a_b") == sm.safe_key("telegram:a:b") + assert sm._storage_key("telegram:a_b") != sm._storage_key("telegram:a:b") def test_save_uses_new_path_not_lossy(tmp_path: Path, monkeypatch) -> None: @@ -93,14 +93,19 @@ def test_load_migrates_lossy_to_new_path(tmp_path: Path, monkeypatch) -> None: assert not lossy_path.exists() -def test_safe_key_is_collision_resistant() -> None: +def test_safe_key_is_lossy() -> None: + assert SessionManager.safe_key("telegram:a_b") == SessionManager.safe_key("telegram:a:b") + + +def test_storage_key_is_collision_resistant() -> None: encoded = { - SessionManager.safe_key("a:b"), - SessionManager.safe_key("a_b"), - SessionManager.safe_key("a:b:c"), + SessionManager._storage_key("a:b"), + SessionManager._storage_key("a_b"), + SessionManager._storage_key("a:b:c"), } assert len(encoded) == 3 + assert SessionManager._storage_key("telegram:a_b") != SessionManager._storage_key("telegram:a:b") def test_lossy_path_helper_returns_expected_path(tmp_path: Path, monkeypatch) -> None: @@ -109,3 +114,32 @@ def test_lossy_path_helper_returns_expected_path(tmp_path: Path, monkeypatch) -> expected = sm.sessions_dir / f"{safe_filename(key.replace(':', '_'))}.jsonl" assert sm._get_legacy_lossy_path(key) == expected + + +def test_storage_paths_are_distinct_when_keys_collide_under_safe_key( + tmp_path: Path, + monkeypatch, +) -> None: + sm = _manager(tmp_path, monkeypatch) + first = Session(key="telegram:a_b") + first.add_message("user", "underscore history") + second = Session(key="telegram:a:b") + second.add_message("user", "colon history") + + sm.save(first) + sm.save(second) + + assert sm.safe_key(first.key) == sm.safe_key(second.key) + assert sm._get_session_path(first.key).exists() + assert sm._get_session_path(second.key).exists() + assert sm._get_session_path(first.key) != sm._get_session_path(second.key) + + sm.invalidate(first.key) + sm.invalidate(second.key) + loaded_first = sm._load(first.key) + loaded_second = sm._load(second.key) + + assert loaded_first is not None + assert loaded_second is not None + assert loaded_first.messages[0]["content"] == "underscore history" + assert loaded_second.messages[0]["content"] == "colon history" diff --git a/tests/agent/test_session_delete.py b/tests/agent/test_session_delete.py index e262cbfc..97955fc1 100644 --- a/tests/agent/test_session_delete.py +++ b/tests/agent/test_session_delete.py @@ -58,11 +58,11 @@ def test_read_session_file_missing(tmp_path: Path) -> None: assert sm.read_session_file("nope:none") is None -def test_safe_key_matches_internal_path(tmp_path: Path) -> None: +def test_storage_key_matches_internal_path(tmp_path: Path) -> None: sm = SessionManager(tmp_path) key = "telegram:abc/def" expected = sm._get_session_path(key).name - assert SessionManager.safe_key(key) + ".jsonl" == expected + assert SessionManager._storage_key(key) + ".jsonl" == expected def _write_legacy_session(legacy_dir: Path, key: str, roles: list[str]) -> Path: