fix(guardrails): bound per_session dedup keys with a fixed-size hash

The dedup cache keyed sessions by the serialized caller, session id and
side, so a caller sending large distinct litellm_session_id values grew it
by their full size for an hour. Keys are now the SHA-256 of that tuple,
64 characters each, so the 100k-entry cap bounds the cache's memory
This commit is contained in:
Caduri Katzav 2026-10-01 18:32:29 +03:00
parent 045b13fe21
commit a37c292b5f
2 changed files with 22 additions and 1 deletions

View file

@ -1,3 +1,4 @@
import hashlib
import json
from typing import Final, Literal, NamedTuple
@ -49,6 +50,10 @@ def _jsonable(value: object) -> object:
return to_jsonable_python(value, fallback=repr, bytes_mode="base64")
def _session_key(caller: Caller, session_id: str, input_type: Literal["request", "response"]) -> str:
return hashlib.sha256(json.dumps((*caller, session_id, input_type)).encode()).hexdigest()
def _sent_value(sent: GenericGuardrailAPIInputs, key: str) -> object:
return sent.get(key) if key in sent else ()
@ -87,7 +92,7 @@ class RecordScope:
case "off":
return False
case "per_session":
return session_id is None or self._claim_session(json.dumps((*caller, session_id, input_type)))
return session_id is None or self._claim_session(_session_key(caller, session_id, input_type))
return assert_never(self._scope)
def _claim_session(self, key: str) -> bool:

View file

@ -96,3 +96,19 @@ def test_per_call_and_off_decide_without_claiming_a_session(scope: GuardrailInfo
)
assert (decisions, session_still_unclaimed) == ([recorded, recorded], True)
def test_a_huge_session_id_is_stored_under_a_fixed_size_key() -> None:
sessions: Final = InMemoryCache()
record_scope: Final = RecordScope("per_session", recorded_sessions=sessions)
caller: Final = Caller(key_hash="hash-a", team_id=None, user_id=None)
huge_session_id: Final = "s" * 1_000_000
decisions: Final = [
record_scope.should_record_allow(session_id=huge_session_id, caller=caller, input_type="request")
for _ in range(2)
]
assert (decisions, [len(key) for key in sessions.cache_dict]) == ([True, False], [64]), (
"a caller-chosen session id must not grow the in-memory dedup cache"
)