fix: prevent kickoff hanging on LLM auth failures

- Move iteration counter increment before LLM call
- Add proper error handling for LLM failures
- Implement max iterations check and handling
- Add error messaging for failed LLM calls

This prevents the agent from hanging in an infinite loop when LLM
authentication fails by properly incrementing the iteration counter
and respecting the max_iterations limit.

Co-Authored-By: brandon@crewai.com <brandon@crewai.com>

src/crewai/agents/crew_agent_executor.py

@@ -108,14 +108,42 @@ class CrewAgentExecutor(CrewAgentExecutorMixin):
         self._create_long_term_memory(formatted_answer)
         return {"output": formatted_answer.output}
 
+    def _has_reached_max_iterations(self) -> bool:
+        """Check if the maximum number of iterations has been reached."""
+        return self.iterations >= self.max_iter
+
+    def _handle_max_iterations_exceeded(self, formatted_answer) -> AgentFinish:
+        """Handle the case when maximum iterations have been exceeded."""
+        error_msg = f"Agent has exceeded maximum iterations ({self.max_iter}). This could be due to repeated failures or infinite loops."
+        self._printer.print(content=error_msg, color="red")
+        return AgentFinish(
+            thought="Maximum iterations exceeded",
+            output=error_msg,
+            text=formatted_answer.text if formatted_answer else error_msg
+        )
+
     def _invoke_loop(self, formatted_answer=None):
         try:
             while not isinstance(formatted_answer, AgentFinish):
                 if not self.request_within_rpm_limit or self.request_within_rpm_limit():
-                    answer = self.llm.call(
-                        self.messages,
-                        callbacks=self.callbacks,
-                    )
+                    # Increment iteration counter before LLM call to prevent hanging
+                    self.iterations += 1
+                    
+                    try:
+                        answer = self.llm.call(
+                            self.messages,
+                            callbacks=self.callbacks,
+                        )
+                    except Exception as llm_error:
+                        self._printer.print(
+                            content=f"LLM call failed: {str(llm_error)}",
+                            color="red",
+                        )
+                        # Check if we've reached max iterations to avoid infinite loops
+                        if self._has_reached_max_iterations():
+                            return self._handle_max_iterations_exceeded(formatted_answer)
+                        # Re-raise authentication and other critical errors
+                        raise
 
                     if answer is None or answer == "":
                         self._printer.print(
@@ -136,7 +164,6 @@ class CrewAgentExecutor(CrewAgentExecutorMixin):
                             ):
                                 answer = answer.split("Observation:")[0].strip()
 
-                    self.iterations += 1
                     formatted_answer = self._format_answer(answer)
 
                     if isinstance(formatted_answer, AgentAction):

src/crewai/memory/entity/entity_memory.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 from crewai.memory.entity.entity_memory_item import EntityMemoryItem
 from crewai.memory.memory import Memory
 from crewai.memory.storage.rag_storage import RAGStorage
@@ -40,7 +38,7 @@ class EntityMemory(Memory):
             )
         super().__init__(storage)
 
-    def save(self, item: EntityMemoryItem, custom_key: Optional[str] = None) -> None:  # type: ignore # BUG?: Signature of "save" incompatible with supertype "Memory"
+    def save(self, item: EntityMemoryItem) -> None:  # type: ignore # BUG?: Signature of "save" incompatible with supertype "Memory"
         """Saves an entity item into the SQLite storage."""
         if self.memory_provider == "mem0":
             data = f"""
@@ -51,7 +49,7 @@ class EntityMemory(Memory):
             """
         else:
             data = f"{item.name}({item.type}): {item.description}"
-        super().save(data, item.metadata, custom_key=custom_key)
+        super().save(data, item.metadata)
 
     def reset(self) -> None:
         try:

src/crewai/memory/long_term/long_term_memory.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List
 
 from crewai.memory.long_term.long_term_memory_item import LongTermMemoryItem
 from crewai.memory.memory import Memory
@@ -19,12 +19,9 @@ class LongTermMemory(Memory):
             storage = LTMSQLiteStorage(db_path=path) if path else LTMSQLiteStorage()
         super().__init__(storage)
 
-    def save(self, item: LongTermMemoryItem, custom_key: Optional[str] = None) -> None:  # type: ignore # BUG?: Signature of "save" incompatible with supertype "Memory"
+    def save(self, item: LongTermMemoryItem) -> None:  # type: ignore # BUG?: Signature of "save" incompatible with supertype "Memory"
         metadata = item.metadata
         metadata.update({"agent": item.agent, "expected_output": item.expected_output})
-        if custom_key:
-            metadata.update({"custom_key": custom_key})
-            
         self.storage.save(  # type: ignore # BUG?: Unexpected keyword argument "task_description","score","datetime" for "save" of "Storage"
             task_description=item.task,
             score=metadata["quality"],
@@ -32,8 +29,8 @@ class LongTermMemory(Memory):
             datetime=item.datetime,
         )
 
-    def search(self, task: str, latest_n: int = 3, custom_key: Optional[str] = None) -> List[Dict[str, Any]]:  # type: ignore # signature of "search" incompatible with supertype "Memory"
-        return self.storage.load(task, latest_n, custom_key)  # type: ignore # BUG?: "Storage" has no attribute "load"
+    def search(self, task: str, latest_n: int = 3) -> List[Dict[str, Any]]:  # type: ignore # signature of "search" incompatible with supertype "Memory"
+        return self.storage.load(task, latest_n)  # type: ignore # BUG?: "Storage" has no attribute "load"
 
     def reset(self) -> None:
         self.storage.reset()

src/crewai/memory/memory.py

@@ -5,10 +5,7 @@ from crewai.memory.storage.rag_storage import RAGStorage
 
 class Memory:
     """
-    Base class for memory, now supporting agent tags, generic metadata, and custom keys.
-    
-    Custom keys allow scoping memories to specific entities (users, accounts, sessions),
-    retrieving memories contextually, and preventing data leakage across logical boundaries.
+    Base class for memory, now supporting agent tags and generic metadata.
     """
 
     def __init__(self, storage: RAGStorage):
@@ -19,13 +16,10 @@ class Memory:
         value: Any,
         metadata: Optional[Dict[str, Any]] = None,
         agent: Optional[str] = None,
-        custom_key: Optional[str] = None,
     ) -> None:
         metadata = metadata or {}
         if agent:
             metadata["agent"] = agent
-        if custom_key:
-            metadata["custom_key"] = custom_key
 
         self.storage.save(value, metadata)
 
@@ -34,12 +28,7 @@ class Memory:
         query: str,
         limit: int = 3,
         score_threshold: float = 0.35,
-        custom_key: Optional[str] = None,
     ) -> List[Any]:
-        filter_dict = None
-        if custom_key:
-            filter_dict = {"custom_key": {"$eq": custom_key}}
-            
         return self.storage.search(
-            query=query, limit=limit, score_threshold=score_threshold, filter=filter_dict
+            query=query, limit=limit, score_threshold=score_threshold
         )

src/crewai/memory/short_term/short_term_memory.py

@@ -46,31 +46,22 @@ class ShortTermMemory(Memory):
         value: Any,
         metadata: Optional[Dict[str, Any]] = None,
         agent: Optional[str] = None,
-        custom_key: Optional[str] = None,
     ) -> None:
         item = ShortTermMemoryItem(data=value, metadata=metadata, agent=agent)
         if self.memory_provider == "mem0":
             item.data = f"Remember the following insights from Agent run: {item.data}"
 
-        super().save(value=item.data, metadata=item.metadata, agent=item.agent, custom_key=custom_key)
+        super().save(value=item.data, metadata=item.metadata, agent=item.agent)
 
     def search(
         self,
         query: str,
         limit: int = 3,
         score_threshold: float = 0.35,
-        custom_key: Optional[str] = None,
     ):
-        filter_dict = None
-        if custom_key:
-            filter_dict = {"custom_key": {"$eq": custom_key}}
-            
         return self.storage.search(
-            query=query, 
-            limit=limit, 
-            score_threshold=score_threshold,
-            filter=filter_dict
-        )
+            query=query, limit=limit, score_threshold=score_threshold
+        )  # type: ignore # BUG? The reference is to the parent class, but the parent class does not have this parameters
 
     def reset(self) -> None:
         try:

src/crewai/memory/storage/ltm_sqlite_storage.py

@@ -70,31 +70,22 @@ class LTMSQLiteStorage:
             )
 
     def load(
-        self, task_description: str, latest_n: int, custom_key: Optional[str] = None
+        self, task_description: str, latest_n: int
     ) -> Optional[List[Dict[str, Any]]]:
         """Queries the LTM table by task description with error handling."""
         try:
             with sqlite3.connect(self.db_path) as conn:
                 cursor = conn.cursor()
-                
-                query = """
+                cursor.execute(
+                    f"""
                     SELECT metadata, datetime, score
                     FROM long_term_memories
                     WHERE task_description = ?
-                """
-                
-                params = [task_description]
-                
-                if custom_key:
-                    query += " AND json_extract(metadata, '$.custom_key') = ?"
-                    params.append(custom_key)
-                
-                query += f"""
                     ORDER BY datetime DESC, score ASC
                     LIMIT {latest_n}
-                """
-                
-                cursor.execute(query, params)
+                """,  # nosec
+                    (task_description,),
+                )
                 rows = cursor.fetchall()
                 if rows:
                     return [

src/crewai/memory/storage/rag_storage.py

@@ -120,11 +120,7 @@ class RAGStorage(BaseRAGStorage):
 
         try:
             with suppress_logging():
-                response = self.collection.query(
-                    query_texts=query, 
-                    n_results=limit,
-                    where=filter
-                )
+                response = self.collection.query(query_texts=query, n_results=limit)
 
             results = []
             for i in range(len(response["ids"][0])):

src/crewai/memory/user/user_memory.py

@@ -26,27 +26,20 @@ class UserMemory(Memory):
         value,
         metadata: Optional[Dict[str, Any]] = None,
         agent: Optional[str] = None,
-        custom_key: Optional[str] = None,
     ) -> None:
         # TODO: Change this function since we want to take care of the case where we save memories for the usr
         data = f"Remember the details about the user: {value}"
-        super().save(data, metadata, custom_key=custom_key)
+        super().save(data, metadata)
 
     def search(
         self,
         query: str,
         limit: int = 3,
         score_threshold: float = 0.35,
-        custom_key: Optional[str] = None,
     ):
-        filter_dict = None
-        if custom_key:
-            filter_dict = {"custom_key": {"$eq": custom_key}}
-            
         results = self.storage.search(
             query=query,
             limit=limit,
             score_threshold=score_threshold,
-            filter=filter_dict,
         )
         return results

tests/memory/custom_key_memory_test.py

@@ -1,57 +0,0 @@
-import pytest
-from unittest.mock import patch, MagicMock
-
-from crewai.memory.short_term.short_term_memory import ShortTermMemory
-from crewai.memory.short_term.short_term_memory_item import ShortTermMemoryItem
-from crewai.agent import Agent
-from crewai.crew import Crew
-from crewai.task import Task
-
-
-@pytest.fixture
-def short_term_memory():
-    """Fixture to create a ShortTermMemory instance"""
-    agent = Agent(
-        role="Researcher",
-        goal="Search relevant data and provide results",
-        backstory="You are a researcher at a leading tech think tank.",
-        tools=[],
-        verbose=True,
-    )
-
-    task = Task(
-        description="Perform a search on specific topics.",
-        expected_output="A list of relevant URLs based on the search query.",
-        agent=agent,
-    )
-    return ShortTermMemory(crew=Crew(agents=[agent], tasks=[task]))
-
-
-def test_save_with_custom_key(short_term_memory):
-    """Test that save method correctly passes custom_key to storage"""
-    with patch.object(short_term_memory.storage, 'save') as mock_save:
-        short_term_memory.save(
-            value="Test data",
-            metadata={"task": "test_task"},
-            agent="test_agent",
-            custom_key="user123",
-        )
-        
-        called_args = mock_save.call_args[0]
-        called_kwargs = mock_save.call_args[1]
-        
-        assert "custom_key" in called_args[1]
-        assert called_args[1]["custom_key"] == "user123"
-
-
-def test_search_with_custom_key(short_term_memory):
-    """Test that search method correctly passes custom_key to storage"""
-    expected_results = [{"context": "Test data", "metadata": {"custom_key": "user123"}, "score": 0.95}]
-    
-    with patch.object(short_term_memory.storage, 'search', return_value=expected_results) as mock_search:
-        results = short_term_memory.search("test query", custom_key="user123")
-        
-        mock_search.assert_called_once()
-        filter_arg = mock_search.call_args[1].get('filter')
-        assert filter_arg == {"custom_key": {"$eq": "user123"}}
-        assert results == expected_results