From 25740710f8e9466d6fbe1917a69011aa17abdf5e Mon Sep 17 00:00:00 2001 From: Shivam Sharma <192242476+ShivamSharma43@users.noreply.github.com> Date: Sun, 11 Oct 2026 20:52:58 +0530 Subject: [PATCH 1/3] Python: Reject oversized MessagePack history records before writing FileHistoryProvider(serialization_format="msgpack") appended records larger than the 64 MiB limit that its reader enforces, so one oversized message made the session unreadable, and every later save failed too because it reads the history for deduplication first. Encode and check the whole batch before opening the file, so a rejected batch leaves the existing history untouched and the call can be retried. Fixes #8633 Co-Authored-By: Claude Opus 5.5 --- .../core/agent_framework/_sessions.py | 16 ++++++--- .../packages/core/tests/core/test_sessions.py | 36 +++++++++++++++++++ 2 files changed, 48 insertions(+), 4 deletions(-) diff --git a/python/packages/core/agent_framework/_sessions.py b/python/packages/core/agent_framework/_sessions.py index 73bfbaa603b..c63783e1e95 100644 --- a/python/packages/core/agent_framework/_sessions.py +++ b/python/packages/core/agent_framework/_sessions.py @@ -2466,11 +2466,19 @@ def _append_messages() -> None: new_messages = filter_new_messages(existing_messages, messages) if not new_messages: return + # Encode and check the whole batch first so a record the reader would reject + # never reaches the file and the existing history stays loadable. + records = [_DEFAULT_MSGPACK_ENCODER.encode(message.to_dict()) for message in new_messages] + for record in records: + if len(record) > self._MAX_MSGPACK_RECORD_BYTES: + raise ValueError( + f"Message encodes to {len(record)} bytes, which exceeds the " + f"{self._MAX_MSGPACK_RECORD_BYTES}-byte MessagePack history record limit." + ) with file_path.open("ab") as file_handle: - for message in new_messages: - serialized = _DEFAULT_MSGPACK_ENCODER.encode(message.to_dict()) - file_handle.write(len(serialized).to_bytes(self._MSGPACK_RECORD_HEADER_BYTES, "big")) - file_handle.write(serialized) + for record in records: + file_handle.write(len(record).to_bytes(self._MSGPACK_RECORD_HEADER_BYTES, "big")) + file_handle.write(record) async with async_lock: await asyncio.to_thread(_append_messages) diff --git a/python/packages/core/tests/core/test_sessions.py b/python/packages/core/tests/core/test_sessions.py index 1cf5d5adcfc..45fd79a65a1 100644 --- a/python/packages/core/tests/core/test_sessions.py +++ b/python/packages/core/tests/core/test_sessions.py @@ -1979,6 +1979,42 @@ async def test_truncated_msgpack_record_raises(self, tmp_path: Path) -> None: with pytest.raises(ValueError, match="record 1.*truncated"): await provider.get_messages("broken") + async def test_msgpack_rejects_oversized_batch_without_mutating_history( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(FileHistoryProvider, "_MAX_MSGPACK_RECORD_BYTES", 256) + provider = FileHistoryProvider(tmp_path, serialization_format="msgpack") + session_file = provider._session_file_path("oversized") + await provider.save_messages("oversized", [Message(role="user", contents=["hello"])]) + saved_bytes = await asyncio.to_thread(session_file.read_bytes) + + with pytest.raises(ValueError, match="exceeds the 256-byte MessagePack history record limit"): + await provider.save_messages( + "oversized", + [ + Message(role="assistant", contents=["fits"]), + Message(role="assistant", contents=["x" * 256]), + ], + ) + + assert await asyncio.to_thread(session_file.read_bytes) == saved_bytes + assert [message.text for message in await provider.get_messages("oversized")] == ["hello"] + + await provider.save_messages("oversized", [Message(role="assistant", contents=["fits"])]) + assert [message.text for message in await provider.get_messages("oversized")] == ["hello", "fits"] + + async def test_msgpack_record_at_size_limit_round_trips( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + message = Message(role="assistant", contents=["at the limit"]) + record_length = len(msgspec.msgpack.encode(message.to_dict())) + monkeypatch.setattr(FileHistoryProvider, "_MAX_MSGPACK_RECORD_BYTES", record_length) + provider = FileHistoryProvider(tmp_path, serialization_format="msgpack") + + await provider.save_messages("limit", [message]) + + assert [loaded.text for loaded in await provider.get_messages("limit")] == ["at the limit"] + async def test_missing_session_file_returns_empty_messages(self, tmp_path: Path) -> None: provider = FileHistoryProvider(tmp_path) From 73cd27c751ae3cf94b273385b4397195f6c2ed50 Mon Sep 17 00:00:00 2001 From: Shivam Sharma <192242476+ShivamSharma43@users.noreply.github.com> Date: Mon, 12 Oct 2026 15:27:30 +0530 Subject: [PATCH 2/3] Python: Bound MessagePack batch preflight memory to one record The size preflight kept every encoded record of the batch, so memory grew with batch length. Encode each record into one reused buffer, once to check the size and once to write, so memory is bounded by the largest record. The batch is still rejected as a whole before the file is opened. Co-Authored-By: Claude Opus 5.5 --- .../packages/core/agent_framework/_sessions.py | 14 +++++++++----- .../packages/core/tests/core/test_sessions.py | 17 +++++++++++++++++ 2 files changed, 26 insertions(+), 5 deletions(-) diff --git a/python/packages/core/agent_framework/_sessions.py b/python/packages/core/agent_framework/_sessions.py index c63783e1e95..55220ad2744 100644 --- a/python/packages/core/agent_framework/_sessions.py +++ b/python/packages/core/agent_framework/_sessions.py @@ -2466,17 +2466,21 @@ def _append_messages() -> None: new_messages = filter_new_messages(existing_messages, messages) if not new_messages: return - # Encode and check the whole batch first so a record the reader would reject - # never reaches the file and the existing history stays loadable. - records = [_DEFAULT_MSGPACK_ENCODER.encode(message.to_dict()) for message in new_messages] - for record in records: + # Check the whole batch first so a record the reader would reject never reaches + # the file and the existing history stays loadable. Records are encoded into one + # reused buffer, once to check and once to write, so memory is bounded by the + # largest record rather than the batch. + record = bytearray() + for message in new_messages: + _DEFAULT_MSGPACK_ENCODER.encode_into(message.to_dict(), record) if len(record) > self._MAX_MSGPACK_RECORD_BYTES: raise ValueError( f"Message encodes to {len(record)} bytes, which exceeds the " f"{self._MAX_MSGPACK_RECORD_BYTES}-byte MessagePack history record limit." ) with file_path.open("ab") as file_handle: - for record in records: + for message in new_messages: + _DEFAULT_MSGPACK_ENCODER.encode_into(message.to_dict(), record) file_handle.write(len(record).to_bytes(self._MSGPACK_RECORD_HEADER_BYTES, "big")) file_handle.write(record) diff --git a/python/packages/core/tests/core/test_sessions.py b/python/packages/core/tests/core/test_sessions.py index 45fd79a65a1..d26638b1d3d 100644 --- a/python/packages/core/tests/core/test_sessions.py +++ b/python/packages/core/tests/core/test_sessions.py @@ -5,6 +5,7 @@ import math import threading import time +import tracemalloc from collections.abc import Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass from pathlib import Path @@ -2015,6 +2016,22 @@ async def test_msgpack_record_at_size_limit_round_trips( assert [loaded.text for loaded in await provider.get_messages("limit")] == ["at the limit"] + async def test_msgpack_batch_save_memory_does_not_grow_with_batch_length(self, tmp_path: Path) -> None: + record_text_bytes = 1024 * 1024 + batch = [Message(role="assistant", contents=[str(index) * record_text_bytes]) for index in range(16)] + provider = FileHistoryProvider(tmp_path, serialization_format="msgpack") + + tracemalloc.start() + try: + await provider.save_messages("large-batch", batch) + _, peak_bytes = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + + # Holding every encoded record at once would need about 16 MiB; one record at a time stays near 1 MiB. + assert peak_bytes < 6 * record_text_bytes + assert len(await provider.get_messages("large-batch")) == len(batch) + async def test_missing_session_file_returns_empty_messages(self, tmp_path: Path) -> None: provider = FileHistoryProvider(tmp_path) From 46e1ce2f5efe62317c7f847f93be2f43ef9759e3 Mon Sep 17 00:00:00 2001 From: Shivam Sharma <192242476+ShivamSharma43@users.noreply.github.com> Date: Thu, 15 Oct 2026 13:44:03 +0530 Subject: [PATCH 3/3] Python: Validate MessagePack records again when writing and roll back on failure The preflight checked the first encoding of each message, but the write pass encoded it again and wrote that result unchecked. A message whose to_dict() changes between the passes could still append a record the reader rejects. The write pass now checks each record before writing it, and any failure truncates the file back to where this call started appending, so the history stays loadable. Co-Authored-By: Claude Opus 5.5 --- .../core/agent_framework/_sessions.py | 32 +++++++++++----- .../packages/core/tests/core/test_sessions.py | 37 +++++++++++++++++++ 2 files changed, 59 insertions(+), 10 deletions(-) diff --git a/python/packages/core/agent_framework/_sessions.py b/python/packages/core/agent_framework/_sessions.py index 55220ad2744..accf14913ee 100644 --- a/python/packages/core/agent_framework/_sessions.py +++ b/python/packages/core/agent_framework/_sessions.py @@ -2472,21 +2472,33 @@ def _append_messages() -> None: # largest record rather than the batch. record = bytearray() for message in new_messages: - _DEFAULT_MSGPACK_ENCODER.encode_into(message.to_dict(), record) - if len(record) > self._MAX_MSGPACK_RECORD_BYTES: - raise ValueError( - f"Message encodes to {len(record)} bytes, which exceeds the " - f"{self._MAX_MSGPACK_RECORD_BYTES}-byte MessagePack history record limit." - ) + self._encode_msgpack_record(message, record) with file_path.open("ab") as file_handle: - for message in new_messages: - _DEFAULT_MSGPACK_ENCODER.encode_into(message.to_dict(), record) - file_handle.write(len(record).to_bytes(self._MSGPACK_RECORD_HEADER_BYTES, "big")) - file_handle.write(record) + append_offset = file_handle.seek(0, os.SEEK_END) + try: + for message in new_messages: + # Encode and check again: the bytes written must be the checked bytes, + # and to_dict() may not return the same result as during the preflight. + self._encode_msgpack_record(message, record) + file_handle.write(len(record).to_bytes(self._MSGPACK_RECORD_HEADER_BYTES, "big")) + file_handle.write(record) + except BaseException: + # Drop any records this call already appended so the history stays loadable. + file_handle.truncate(append_offset) + raise async with async_lock: await asyncio.to_thread(_append_messages) + def _encode_msgpack_record(self, message: Message, record: bytearray) -> None: + """Encode ``message`` into ``record`` and reject it if the reader would refuse the record.""" + _DEFAULT_MSGPACK_ENCODER.encode_into(message.to_dict(), record) + if len(record) > self._MAX_MSGPACK_RECORD_BYTES: + raise ValueError( + f"Message encodes to {len(record)} bytes, which exceeds the " + f"{self._MAX_MSGPACK_RECORD_BYTES}-byte MessagePack history record limit." + ) + def _read_json_messages(self, file_path: Path) -> list[Message]: """Read JSON Lines messages from ``file_path``.""" messages: list[Message] = [] diff --git a/python/packages/core/tests/core/test_sessions.py b/python/packages/core/tests/core/test_sessions.py index d26638b1d3d..99a1d3cc767 100644 --- a/python/packages/core/tests/core/test_sessions.py +++ b/python/packages/core/tests/core/test_sessions.py @@ -2004,6 +2004,43 @@ async def test_msgpack_rejects_oversized_batch_without_mutating_history( await provider.save_messages("oversized", [Message(role="assistant", contents=["fits"])]) assert [message.text for message in await provider.get_messages("oversized")] == ["hello", "fits"] + async def test_msgpack_rejects_record_that_grows_after_preflight_without_mutating_history( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + class GrowingMessage(Message): + """Fits the record limit on the first to_dict() call and exceeds it afterwards.""" + + to_dict_calls = 0 + + def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]: + GrowingMessage.to_dict_calls += 1 + if GrowingMessage.to_dict_calls > 1: + return Message(role="assistant", contents=["x" * 256]).to_dict() + return super().to_dict(exclude=exclude, exclude_none=exclude_none) + + monkeypatch.setattr(FileHistoryProvider, "_MAX_MSGPACK_RECORD_BYTES", 256) + provider = FileHistoryProvider(tmp_path, serialization_format="msgpack") + session_file = provider._session_file_path("growing") + await provider.save_messages("growing", [Message(role="user", contents=["hello"])]) + saved_bytes = await asyncio.to_thread(session_file.read_bytes) + + # The first message is written before the grown record is found, so the rollback must remove it too. + with pytest.raises(ValueError, match="exceeds the 256-byte MessagePack history record limit"): + await provider.save_messages( + "growing", + [ + Message(role="assistant", contents=["fits"]), + GrowingMessage(role="assistant", contents=["grows"]), + ], + ) + + assert GrowingMessage.to_dict_calls == 2 + assert await asyncio.to_thread(session_file.read_bytes) == saved_bytes + assert [message.text for message in await provider.get_messages("growing")] == ["hello"] + + await provider.save_messages("growing", [Message(role="assistant", contents=["fits"])]) + assert [message.text for message in await provider.get_messages("growing")] == ["hello", "fits"] + async def test_msgpack_record_at_size_limit_round_trips( self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: