From c5e35cf1216ad40dcbde706fe17afbaa7939c748 Mon Sep 17 00:00:00 2001 From: citizen204 Date: Sat, 27 Jun 2026 18:21:22 +0930 Subject: [PATCH] fix(memory): strip injected user context from restored messages to prevent O(turns) accumulation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When retrieval_config is set, retrieve_customer_context() prepends a block to each user message and the message is persisted with that block in storage. On every subsequent turn list_messages() reloads all historical messages verbatim, so the injected block reappears in every restored user message — O(turns × records) tokens. Add _filter_restored_user_context() (mirroring the existing _filter_restored_tool_context pattern) and call it from list_messages() whenever retrieval_config is configured. The filter removes any content block whose text starts with the opening context tag before the messages are handed to the Strands runtime. Fresh context is re-injected on each new turn by retrieve_customer_context(), so no context is lost. Fixes #420 --- .../integrations/strands/session_manager.py | 36 +++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/src/bedrock_agentcore/memory/integrations/strands/session_manager.py b/src/bedrock_agentcore/memory/integrations/strands/session_manager.py index c4fe8540..c4380c7d 100644 --- a/src/bedrock_agentcore/memory/integrations/strands/session_manager.py +++ b/src/bedrock_agentcore/memory/integrations/strands/session_manager.py @@ -787,6 +787,8 @@ def list_messages( messages = self.converter.events_to_messages(events) if self.config.filter_restored_tool_context: messages = self._filter_restored_tool_context(messages) + if self.config.retrieval_config: + messages = self._filter_restored_user_context(messages) if limit is not None: return messages[offset : offset + limit] else: @@ -823,6 +825,40 @@ def _filter_restored_tool_context(self, messages: list[SessionMessage]) -> list[ return filtered_messages + def _filter_restored_user_context(self, messages: list[SessionMessage]) -> list[SessionMessage]: + """Strip previously-injected blocks from restored user messages. + + Without this, the context block that was prepended to each user message during + a prior turn is reloaded from storage on every subsequent turn, causing the + injected context to accumulate O(turns × records) across the conversation. + Fresh context is re-injected by retrieve_customer_context on each new turn. + """ + open_tag = f"<{self.config.context_tag}>" + filtered_messages: list[SessionMessage] = [] + for session_message in messages: + message = session_message.to_message() + if message.get("role") != "user": + filtered_messages.append(session_message) + continue + filtered_content = [ + block + for block in message.get("content", []) + if not (isinstance(block, dict) and block.get("text", "").startswith(open_tag)) + ] + if not filtered_content: + filtered_messages.append(session_message) + continue + filtered_messages.append( + SessionMessage( + message={"role": message["role"], "content": filtered_content}, + message_id=session_message.message_id, + redact_message=session_message.redact_message, + created_at=session_message.created_at, + updated_at=session_message.updated_at, + ) + ) + return filtered_messages + # endregion SessionRepository interface implementation # region RepositorySessionManager overrides