From 61cd1dad52477dcd5f8628af8c47160d5ff63653 Mon Sep 17 00:00:00 2001 From: Padma Komarina Date: Sun, 1 Mar 2026 20:35:48 -0500 Subject: [PATCH] fix: AgentCoreMemorySessionManager - Cache agent timestamps to eliminate redundant list_events calls --- .../integrations/strands/session_manager.py | 25 ++++++++++++++----- .../test_agentcore_memory_session_manager.py | 17 +++++++++++++ 2 files changed, 36 insertions(+), 6 deletions(-) diff --git a/src/bedrock_agentcore/memory/integrations/strands/session_manager.py b/src/bedrock_agentcore/memory/integrations/strands/session_manager.py index a40bcd86..2783a498 100644 --- a/src/bedrock_agentcore/memory/integrations/strands/session_manager.py +++ b/src/bedrock_agentcore/memory/integrations/strands/session_manager.py @@ -122,6 +122,9 @@ def __init__( self._message_buffer: list[tuple[str, list[tuple[str, str]], bool, datetime]] = [] self._buffer_lock = threading.Lock() + # Cache for agent created_at timestamps to avoid fetching on every update + self._agent_created_at_cache: dict[str, datetime] = {} + # Add strands-agents to the request user agent if boto_client_config: existing_user_agent = getattr(boto_client_config, "user_agent_extra", None) @@ -280,6 +283,11 @@ def create_agent(self, session_id: str, session_agent: SessionAgent, **kwargs: A AGENT_ID_KEY: {"stringValue": session_agent.agent_id}, }, ) + + # Cache the created_at timestamp to avoid re-fetching on updates + if session_agent.created_at: + self._agent_created_at_cache[session_agent.agent_id] = session_agent.created_at + logger.info( "Created agent: %s in session: %s with event %s", session_agent.agent_id, @@ -327,7 +335,11 @@ def read_agent(self, session_id: str, agent_id: str, **kwargs: Any) -> Optional[ if events: agent_data = json.loads(events[0].get("payload", {})[0].get("blob")) - return SessionAgent.from_dict(agent_data) + agent = SessionAgent.from_dict(agent_data) + # Cache the created_at timestamp to avoid re-fetching on updates + if agent.created_at: + self._agent_created_at_cache[agent_id] = agent.created_at + return agent # 2. Fallback: check for legacy event and migrate legacy_actor_id = f"{LEGACY_AGENT_PREFIX}{agent_id}" @@ -369,11 +381,12 @@ def update_agent(self, session_id: str, session_agent: SessionAgent, **kwargs: A SessionException: If session ID doesn't match configuration. """ agent_id = session_agent.agent_id - previous_agent = self.read_agent(session_id=session_id, agent_id=agent_id) - if previous_agent is None: - raise SessionException(f"Agent {agent_id} in session {session_id} does not exist") - else: - session_agent.created_at = previous_agent.created_at + + if agent_id not in self._agent_created_at_cache: + previous_agent = self.read_agent(session_id=session_id, agent_id=agent_id) + if previous_agent is None: + raise SessionException(f"Agent {agent_id} in session {session_id} does not exist") + session_agent.created_at = self._agent_created_at_cache[agent_id] # Create a new agent as AgentCore Memory is immutable. We always get the latest one in `read_agent` self.create_agent(session_id, session_agent) diff --git a/tests/bedrock_agentcore/memory/integrations/strands/test_agentcore_memory_session_manager.py b/tests/bedrock_agentcore/memory/integrations/strands/test_agentcore_memory_session_manager.py index 19b3ec23..dd131063 100644 --- a/tests/bedrock_agentcore/memory/integrations/strands/test_agentcore_memory_session_manager.py +++ b/tests/bedrock_agentcore/memory/integrations/strands/test_agentcore_memory_session_manager.py @@ -422,6 +422,23 @@ def test_update_agent(self, session_manager, mock_memory_client): # Should not raise any exceptions session_manager.update_agent("test-session-456", session_agent) + def test_update_agent_uses_cache(self, session_manager, mock_memory_client): + """Test that update_agent uses cache to avoid fetching memory events on subsequent updates.""" + + # Manually populate the cache (simulating what happens after first agent creation/read) + created_at = "2024-01-01T12:00:00+00:00" + session_manager._agent_created_at_cache["test-agent-123"] = created_at + + # Now update the agent - should NOT call list_events since it's in cache + updated_agent = SessionAgent(agent_id="test-agent-123", state={"key": "value"}, conversation_manager_state={}) + session_manager.update_agent("test-session-456", updated_agent) + + # Verify that list_events was NOT called (cache was used) + mock_memory_client.list_events.assert_not_called() + + # Verify the created_at was preserved from cache + assert updated_agent.created_at == created_at + def test_update_agent_wrong_session(self, session_manager): """Test updating an agent with wrong session ID.""" session_agent = SessionAgent(agent_id="test-agent-123", state={}, conversation_manager_state={})