Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 15 additions & 3 deletions src/agents/extensions/memory/encrypt_session.py
Original file line number Diff line number Diff line change
Expand Up @@ -165,10 +165,22 @@ def _unwrap(self, item: TResponseInputItem | EncryptedEnvelope) -> TResponseInpu
return cast(TResponseInputItem, item)

try:
token = item["payload"].encode("utf-8")
payload = item["payload"]
if not isinstance(payload, str):
return None
token = payload.encode("utf-8")
plaintext = self.cipher.decrypt(token, ttl=self.ttl)
return cast(TResponseInputItem, _from_json_bytes(plaintext))
except (InvalidToken, KeyError):
decoded = _from_json_bytes(plaintext)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Handle ValueError from oversized JSON integers

When an authenticated envelope decrypts to a JSON object containing an integer longer than Python's configured digit limit (for example, more than 4,300 digits by default), json.loads() raises ValueError, not JSONDecodeError. That exception escapes here, so get_items() and pop_item() still abort instead of isolating the malformed record as intended; catch this deserialization failure at the same boundary.

AGENTS.md reference: AGENTS.md:L165-L166

Useful? React with 👍 / 👎.

if not isinstance(decoded, dict):
return None
return cast(TResponseInputItem, decoded)
except (
InvalidToken,
KeyError,
UnicodeDecodeError,
UnicodeEncodeError,
json.JSONDecodeError,
):
return None

def _unwrap_valid_items(
Expand Down
48 changes: 48 additions & 0 deletions tests/extensions/memory/test_encrypt_session.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,20 @@ def _invalid_encrypted_envelope() -> TResponseInputItem:
)


def _malformed_encrypted_envelope() -> TResponseInputItem:
return cast(
TResponseInputItem,
{"__enc__": 1, "v": 1, "kid": "hkdf-v1", "payload": None},
)


def _malformed_unicode_encrypted_envelope() -> TResponseInputItem:
return cast(
TResponseInputItem,
{"__enc__": 1, "v": 1, "kid": "hkdf-v1", "payload": "\ud800"},
)


@pytest.fixture
def agent() -> Agent:
"""Fixture for a basic agent with a scripted model."""
Expand Down Expand Up @@ -381,6 +395,40 @@ async def test_encrypted_session_get_items_limit_skips_invalid_latest_envelope(
underlying_session.close()


async def test_encrypted_session_skips_malformed_envelopes(
encryption_key: str, underlying_session: SQLiteSession
):
"""Malformed persisted envelopes should be skipped like invalid tokens."""
session = EncryptedSession(
session_id="test_session",
underlying_session=underlying_session,
encryption_key=encryption_key,
)

await session.add_items([{"role": "user", "content": "valid"}])
await underlying_session.add_items(
[_malformed_encrypted_envelope(), _malformed_unicode_encrypted_envelope()]
)
await underlying_session.add_items(
[
cast(
TResponseInputItem,
{
"__enc__": 1,
"v": 1,
"kid": "hkdf-v1",
"payload": session.cipher.encrypt(b"[]").decode("utf-8"),
},
)
]
)

items = await session.get_items()
assert [item.get("content") for item in items] == ["valid"]

underlying_session.close()


async def test_encrypted_session_get_items_limit_returns_latest_valid_items_after_invalids(
encryption_key: str, underlying_session: SQLiteSession
):
Expand Down