fix(memory): reuse shared credential redaction before persistence

This commit is contained in:
moe-berri 2026-09-14 18:42:53 -07:00
parent f2026051a3
commit 6f024adfcc
3 changed files with 32 additions and 6 deletions

View file

@ -3,21 +3,18 @@ from typing import Final
from rapidfuzz import fuzz, process
from litellm.litellm_core_utils.secret_redaction import redact_string
from litellm.types.memory_v2 import MemoryEntry
_STOP_WORDS: Final = frozenset(
"the and for how what why with this that does have from about our are was when should can you work team".split()
)
_TOKEN: Final = re.compile(r"[\w-]{2,}", re.UNICODE)
_PRIVATE_KEY: Final = re.compile(r"-----BEGIN [A-Z ]*PRIVATE KEY-----.*?-----END [A-Z ]*PRIVATE KEY-----", re.DOTALL)
_CREDENTIAL: Final = re.compile(r"\b(?:sk-|gh[pousr]_|github_pat_)[A-Za-z0-9_-]{12,}")
_BEARER: Final = re.compile(r"(Bearer\s+)[A-Za-z0-9._~+/-]{12,}", re.IGNORECASE)
def redact_memory(value: str) -> str:
return _BEARER.sub(
r"\1[REDACTED]", _CREDENTIAL.sub("[REDACTED TOKEN]", _PRIVATE_KEY.sub("[REDACTED PRIVATE KEY]", value))
)
return _CREDENTIAL.sub("[REDACTED TOKEN]", redact_string(value))
def _similarity(term: str, text: str, words: tuple[str, ...]) -> float:

View file

@ -55,4 +55,4 @@ def test_recognizable_credentials_are_redacted_without_removing_the_observation(
assert "The integration failed" in redacted
assert "abcdefghijklmnopqrst" not in redacted
assert "private material" not in redacted
assert "[REDACTED PRIVATE KEY]" in redacted
assert "REDACTED" in redacted

View file

@ -329,6 +329,35 @@ async def test_capture_records_authenticated_contributor_and_tenant(database: Ma
assert data["organization_id"] == "org" and data["team_id"] == "team"
@pytest.mark.asyncio
@pytest.mark.parametrize(
("credential", "secret"),
(
("AKIA" + "A" * 16, "A" * 16),
("aws_secret_access_key=" + "B" * 40, "B" * 40),
("AIza" + "C" * 35, "C" * 35),
("postgres://test-user:memory-test-password@db.example.test/app", "memory-test-password"),
("Authorization: Basic " + "D" * 24, "D" * 24),
("Authorization: Bearer " + "E" * 24, "E" * 24),
("ghp_" + "F" * 24, "F" * 24),
("github_pat_" + "G" * 24, "G" * 24),
),
)
async def test_capture_redacts_credentials_before_persisting_content_and_metadata(
database: MagicMock, credential: str, secret: str
) -> None:
configure(database)
text = "Deployment uses staging port 8123; credential: " + credential
fields = ("title", "content", "evidence", "when_to_use", "scope", "source")
await management.capture_entry(_CAPTURE.model_copy(update={field: text for field in fields}), auth())
data = database.db.litellm_memorytable.create.call_args.kwargs["data"]
stored = {"content": data["value"], **json.loads(data["metadata"])}
for field in fields:
assert secret not in stored[field]
assert "REDACTED" in stored[field]
assert "staging port 8123" in stored[field]
@pytest.mark.asyncio
async def test_search_finds_an_old_record_beyond_the_first_thousand(database: MagicMock) -> None:
configure(database)