fix(session): split safe_key and _storage_key to fix WebUI coupling (#4533)
This commit is contained in:
@@ -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."""
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user