Python: Reject oversized MessagePack history records before writing

python/packages/core/agent_framework/_sessions.py

@@ -2467,15 +2467,39 @@ def _append_messages() -> None:
                 new_messages = filter_new_messages(existing_messages, messages)
                 if not new_messages:
                     return
+                # 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:
+                    self._encode_msgpack_record(message, record)
                 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)
+                    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] = []

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
@@ -1981,6 +1982,95 @@ 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_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:
+        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_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)