feat(memory): implement true message batching for flush_messages() #256
+1,056
−34
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add optional batch_size parameter to AgentCoreMemoryConfig that allows customers to buffer messages before sending to AgentCore Memory. Changes: - Add batch_size parameter (default=1, max=100) to AgentCoreMemoryConfig - Add message buffering with thread-safe _message_buffer and _buffer_lock - Modify create_message() to buffer when batch_size > 1 - Add flush_messages() to send all buffered messages - Add pending_message_count() to check buffer size - Add close() and context manager for cleanup - Add 24 unit tests covering batching functionality Default batch_size=1 preserves backward compatibility (immediate send).
Previously, batch_size buffered messages but still made N separate API calls for N messages. Now flush_messages() groups conversational messages by session_id and combines them into a single create_event() call per session, significantly reducing API calls. Key changes: - Group conversational messages by session_id into combined payloads - Preserve message order within each session's payload (earliest first) - Use the latest timestamp from grouped messages for the combined event - Send blob messages (>9KB) individually (different API path) - Clear buffer only after ALL API calls succeed to prevent data loss - Improve error messages to include session context Tests added for critical scenarios: - Multiple sessions grouped into separate API calls - Latest timestamp used for combined events - Partial failure with multiple sessions preserves entire buffer - Multiple blob messages sent individually (not batched) - Mixed sessions with blobs and conversational messages Also fixes pre-existing test issues: - Fix test_read_agent_legacy_migration mock setup to match actual impl - Fix test_load_long_term_memories_with_validation_failure for strands API
jariy17
commented
Feb 6, 2026
|
|
||
| # region Batching support | ||
|
|
||
| def flush_messages(self) -> list[dict[str, Any]]: |
Contributor
Author
There was a problem hiding this comment.
this should be private method.
| """ | ||
| self.flush_messages() | ||
|
|
||
| # endregion Batching support |
Contributor
Author
There was a problem hiding this comment.
Why is this a comment?
| Returns: | ||
| int: Number of buffered messages waiting to be sent. | ||
| """ | ||
| with self._buffer_lock: |
Contributor
Author
There was a problem hiding this comment.
There must be a some atomic library that would incorporate the lock whenever you do operations in _message_buffer.
| return [] | ||
|
|
||
| results = [] | ||
| for session_id, messages, is_blob, monotonic_timestamp in messages_to_send: |
Contributor
Author
There was a problem hiding this comment.
this is not actually buffering right now.
Contributor
Author
There was a problem hiding this comment.
you can take all the messages_to_send and put it under messages, if it's not a blob. https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-agentcore/client/create_event.html. Make sure the messages are ordered in the earilers to latest when storing in one create_event.
Extract _create_session_manager helper to eliminate triple-nested context managers and move batching_config/batching_session_manager to module-level fixtures shared across all test classes.
…in __exit__ Wrap _flush_messages() in try/except so that if the with-block raises and flush also fails, the flush error is logged instead of replacing the original exception.
Iterate reversed(events) instead of reversing the flat message list after iteration, which scrambled intra-event payload order for batched events. Add 6 unit tests covering multi-event ordering, batched payload ordering, mixed blob/conversational ordering, and malformed payloads.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary
batch_sizeconfiguration option for message buffering (1-100, default 1)flush_messages()- groups conversational messages by session_id into single API callsevents_to_messages()ordering bug: iteratereversed(events)instead of reversing the flat message list after iteration, which scrambled intra-event payload order for batched eventsTest plan
pytest tests/bedrock_agentcore/memory/integrations/strands/test_agentcore_memory_session_manager.py -k "Batching" -vpytest tests/bedrock_agentcore/memory/integrations/strands/test_bedrock_converter.py -v