Repository navigation
Python: Reject oversized MessagePack history records before writing - #9260
Shivam Sharma (ShivamSharma43) wants to merge 4 commits into
Conversation
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 microsoft#8633 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
GitHub Actions
Active
GitHub Actions
Active
agent-framework-automation
Bot
added
the
python
|
/review |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
MAF Automated Review — Iteration 1
Result: Findings reported
Scope: full PR (1 commit(s)): 25740710f8e9
Model: gpt-5.6-sol
Overview
The change correctly aligns MessagePack writes with the reader's per-record limit and validates the complete batch before opening the file, with tests covering rollback and the exact-limit boundary. Existing locking, deduplication, and error propagation preserve the intended history contract. However, the preflight materializes every encoded record at once, so an unbounded batch can exhaust shared process memory even when every individual record is valid.
Reviewed the supplied pull-request change set across correctness, security/reliability, architecture, and failure behavior.
1 verified finding remained after source verification (1 medium) across 1 file. Details are attached to the affected lines below.
Affected areas: python/packages/core/agent_framework/_sessions.py
| 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] |
There was a problem hiding this comment.
Because this comprehension retains the encoding of every record in an unbounded batch, a caller can repeat an individually valid near-64 MiB message and force memory usage to scale with the whole batch, potentially OOMing a shared worker. Please preserve the all-or-nothing preflight while bounding aggregate retained memory independently of batch length, for example with a bounded-memory validation/spooling strategy.
There was a problem hiding this comment.
Agreed, fixed in 73cd27c. The preflight no longer keeps the encodings. Each message is encoded into one reused bytearray with Encoder.encode_into, once to check its size and once more when writing, so memory is bounded by the largest record rather than the batch. The whole batch is still checked before the file is opened, so a rejected batch leaves the history untouched.
test_msgpack_batch_save_memory_does_not_grow_with_batch_length saves 16 messages of about 1 MiB each and checks the tracemalloc peak stays under 6 MiB. It peaks at about 3.3 MB with the change and 34.7 MB on 2574071. The existing rollback and exact-limit tests still pass.
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 <noreply@anthropic.com>
GitHub Actions
Active
|
/review |
There was a problem hiding this comment.
MAF Automated Review — Iteration 2
Result: Findings reported
Scope: 1 net-new commit(s): 73cd27c751ae
Model: gpt-5.6-sol
Overview
This incremental change replaces batch-wide encoded-record retention with a reusable buffer, and its regression test confirms that memory no longer scales with batch length. The existing preflight, exact-limit test, and reader-side length guard remain strong defenses. However, the write pass serializes each mutable message again without validating the bytes actually persisted, which can recreate the oversized-record corruption this PR is intended to prevent.
Reviewed the supplied incremental change set across correctness, security/reliability, architecture, and failure behavior.
1 verified finding remained after source verification (1 medium) across 1 file. Details are attached to the affected lines below.
Affected areas: python/packages/core/agent_framework/_sessions.py
| 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) | ||
| _DEFAULT_MSGPACK_ENCODER.encode_into(message.to_dict(), record) |
There was a problem hiding this comment.
The preflight validates the first encoding, but this second to_dict() result is the one actually written and its length is never checked. A mutable message or stateful subclass can fit during preflight and exceed the limit here, leaving a record that the reader rejects and making the session unreadable for later loads and saves. Please ensure the bytes written are the same validated bytes, such as through bounded staging, or make this second pass validation rollback-safe so an invalid record cannot remain appended.
There was a problem hiding this comment.
Fixed in 46e1ce2. The write pass now checks each record's length again before writing it, using the same helper as the preflight, so only checked bytes reach the file. If anything fails during the write pass, including a record that grew since the preflight, the file is truncated back to the offset where this call started appending and the error is raised, so a bad record can't stay appended. Memory is still bounded to one reused record buffer.
test_msgpack_rejects_record_that_grows_after_preflight_without_mutating_history uses a message whose to_dict() fits the first time and is over the limit the second time. It comes after a valid message in the batch, so the valid one is written before the failure. The test checks that the file is byte-for-byte unchanged, that history still loads, and that later saves work. On 73cd27c it fails with DID NOT RAISE.
I also merged the latest main (cbe6e6e). The earlier Windows CI failure was the upstream flaky declarative test test_auth_context_binding_survives_workflow_resume[True-True] (packages/declarative/tests/test_invoke_mcp_tool_executor.py:470), which this PR doesn't touch. The other Windows and 3.15 jobs were cancelled after it failed.
… 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 <noreply@anthropic.com>
GitHub Actions
Active
Motivation & Context
FileHistoryProvider(serialization_format="msgpack")can save a record that its own reader refuses to load._read_msgpack_messages()rejects records longer than_MAX_MSGPACK_RECORD_BYTES(64 MiB), butsave_messages()appended them without checking. After one oversized message,get_messages()raisesValueErrorfor that session, and every latersave_messages()fails the same way because it reads the existing history for deduplication first. A single large inline media payload can make an otherwise valid conversation unusable.Description & Review Guide
FileHistoryProvider.save_messages(), the batch of new messages is now encoded up front, and each record is checked against_MAX_MSGPACK_RECORD_BYTESbefore the file is opened. An oversized record raisesValueErrornaming its size and the limit. Records are then written from the already-encoded bytes, so each message is encoded once.> limitcheck. JSON Lines storage and successful MessagePack saves are unchanged.ValueErrorfromsave_messages()is the right signal here. It matches whatget_messages()already raises for the same record.Tests added in
tests/core/test_sessions.py, using a lowered limit so no large payloads are allocated:test_msgpack_rejects_oversized_batch_without_mutating_history: a batch of a valid record followed by an oversized one is rejected, the file bytes are unchanged, the session still loads, and saving the valid record afterwards works. It fails without the fix.test_msgpack_record_at_size_limit_round_trips: a record exactly at the limit saves and loads.Related Issue
Fixes #8633
Contribution Checklist
breaking changelabel (or add "[BREAKING]" to the title prefix, before or after any language prefix) — a workflow keeps the label and title prefix in sync automatically.🤖 Generated with Claude Code