diff --git a/openhands/controller/agent_controller.py b/openhands/controller/agent_controller.py index 586386bd71..3d3e611140 100644 --- a/openhands/controller/agent_controller.py +++ b/openhands/controller/agent_controller.py @@ -93,6 +93,7 @@ class AgentController: ChangeAgentStateAction, AgentStateChangedObservation, ) + _cached_first_user_message: MessageAction | None = None def __init__( self, @@ -1206,15 +1207,19 @@ class AgentController: Returns: MessageAction | None: The first user message, or None if no user message found """ - # Find the first user message from the appropriate starting point - user_messages = list(self.event_stream.get_events(start_id=self.state.start_id)) + # Return cached message if any + if self._cached_first_user_message is not None: + return self._cached_first_user_message - # Get and return the first user message - return next( + # Find the first user message + self._cached_first_user_message = next( ( e - for e in user_messages + for e in self.event_stream.get_events( + start_id=self.state.start_id, + ) if isinstance(e, MessageAction) and e.source == EventSource.USER ), None, ) + return self._cached_first_user_message diff --git a/tests/unit/test_agent_controller.py b/tests/unit/test_agent_controller.py index fc59069bc9..36b45363bc 100644 --- a/tests/unit/test_agent_controller.py +++ b/tests/unit/test_agent_controller.py @@ -993,6 +993,7 @@ async def test_first_user_message_with_identical_content(): """ Test that _first_user_message correctly identifies the first user message even when multiple messages have identical content but different IDs. + Also verifies that the result is properly cached. The issue we're checking is that the comparison (action == self._first_user_message()) should correctly differentiate between messages with the same content but different IDs. @@ -1038,6 +1039,20 @@ async def test_first_user_message_with_identical_content(): second_message.id != first_user_message.id ) # This should be False, but may be True if there's a bug + # Verify caching behavior + assert ( + controller._cached_first_user_message is not None + ) # Cache should be populated + assert ( + controller._cached_first_user_message is first_user_message + ) # Cache should store the same object + + # Mock get_events to verify it's not called again + with patch.object(event_stream, 'get_events') as mock_get_events: + cached_message = controller._first_user_message() + assert cached_message is first_user_message # Should return cached object + mock_get_events.assert_not_called() # Should not call get_events again + await controller.close()