🐛 Use conversation_id and user_id as the primary key of the agent run thread pool

Author: Phinease

backend/agents/agent_run_manager.py

@@ -21,39 +21,46 @@ def __new__(cls):
 
     def __init__(self):
         if not self._initialized:
-            # conversation_id -> agent_run_info
-            self.agent_runs: Dict[int, AgentRunInfo] = {}
+            # user_id:conversation_id -> agent_run_info
+            self.agent_runs: Dict[str, AgentRunInfo] = {}
             self._initialized = True
 
-    def register_agent_run(self, conversation_id: int, agent_run_info):
+    def _get_run_key(self, conversation_id: int, user_id: str) -> str:
+        """Generate unique key for agent run using user_id and conversation_id"""
+        return f"{user_id}:{conversation_id}"
+
+    def register_agent_run(self, conversation_id: int, agent_run_info, user_id: str):
         """register agent run instance"""
         with self._lock:
-            self.agent_runs[conversation_id] = agent_run_info
+            run_key = self._get_run_key(conversation_id, user_id)
+            self.agent_runs[run_key] = agent_run_info
             logger.info(
-                f"register agent run instance, conversation_id: {conversation_id}")
+                f"register agent run instance, user_id: {user_id}, conversation_id: {conversation_id}")
 
-    def unregister_agent_run(self, conversation_id: int):
+    def unregister_agent_run(self, conversation_id: int, user_id: str):
         """unregister agent run instance"""
         with self._lock:
-            if conversation_id in self.agent_runs:
-                del self.agent_runs[conversation_id]
+            run_key = self._get_run_key(conversation_id, user_id)
+            if run_key in self.agent_runs:
+                del self.agent_runs[run_key]
                 logger.info(
-                    f"unregister agent run instance, conversation_id: {conversation_id}")
+                    f"unregister agent run instance, user_id: {user_id}, conversation_id: {conversation_id}")
             else:
                 logger.info(
-                    f"no agent run instance found for conversation_id: {conversation_id}")
+                    f"no agent run instance found for user_id: {user_id}, conversation_id: {conversation_id}")
 
-    def get_agent_run_info(self, conversation_id: int):
+    def get_agent_run_info(self, conversation_id: int, user_id: str):
         """get agent run instance"""
-        return self.agent_runs.get(conversation_id)
+        run_key = self._get_run_key(conversation_id, user_id)
+        return self.agent_runs.get(run_key)
 
-    def stop_agent_run(self, conversation_id: int) -> bool:
-        """stop agent run for specified conversation_id"""
-        agent_run_info = self.get_agent_run_info(conversation_id)
+    def stop_agent_run(self, conversation_id: int, user_id: str) -> bool:
+        """stop agent run for specified conversation_id and user_id"""
+        agent_run_info = self.get_agent_run_info(conversation_id, user_id)
         if agent_run_info is not None:
             agent_run_info.stop_event.set()
             logger.info(
-                f"agent run stopped, conversation_id: {conversation_id}")
+                f"agent run stopped, user_id: {user_id}, conversation_id: {conversation_id}")
             return True
         return False
 

backend/apps/agent_app.py

@@ -45,11 +45,12 @@ async def agent_run_api(agent_request: AgentRequest, http_request: Request, auth
 
 
 @router.get("/stop/{conversation_id}")
-async def agent_stop_api(conversation_id: int):
+async def agent_stop_api(conversation_id: int, authorization: Optional[str] = Header(None)):
     """
     stop agent run and preprocess tasks for specified conversation_id
     """
-    if stop_agent_tasks(conversation_id).get("status") == "success":
+    user_id, _ = get_current_user_id(authorization)
+    if stop_agent_tasks(conversation_id, user_id).get("status") == "success":
         return {"status": "success", "message": "agent run and preprocess tasks stopped successfully"}
     else:
         raise HTTPException(status_code=HTTPStatus.BAD_REQUEST,

backend/services/agent_service.py

@@ -125,7 +125,7 @@ async def _stream_agent_chunks(
                 user_id=user_id,
             )
         # Always unregister the run to release resources
-        agent_run_manager.unregister_agent_run(agent_request.conversation_id)
+        agent_run_manager.unregister_agent_run(agent_request.conversation_id, user_id)
 
         # Schedule memory addition in background to avoid blocking SSE termination
         async def _add_memory_background():
@@ -681,7 +681,7 @@ async def prepare_agent_run(
         allow_memory_search=allow_memory_search,
     )
     agent_run_manager.register_agent_run(
-        agent_request.conversation_id, agent_run_info)
+        agent_request.conversation_id, agent_run_info, user_id)
     return agent_run_info, memory_context
 
 
@@ -885,13 +885,13 @@ async def run_agent_stream(
     )
 
 
-def stop_agent_tasks(conversation_id: int):
+def stop_agent_tasks(conversation_id: int, user_id: str):
     """
     Stop agent run and preprocess tasks for the specified conversation_id.
     Matches the behavior of agent_app.agent_stop_api.
     """
     # Stop agent run
-    agent_stopped = agent_run_manager.stop_agent_run(conversation_id)
+    agent_stopped = agent_run_manager.stop_agent_run(conversation_id, user_id)
 
     # Stop preprocess tasks
     preprocess_stopped = preprocess_manager.stop_preprocess_tasks(
@@ -904,11 +904,11 @@ def stop_agent_tasks(conversation_id: int):
         if preprocess_stopped:
             message_parts.append("preprocess tasks")
 
-        message = f"successfully stopped {' and '.join(message_parts)} for conversation_id {conversation_id}"
+        message = f"successfully stopped {' and '.join(message_parts)} for user_id {user_id}, conversation_id {conversation_id}"
         logging.info(message)
         return {"status": "success", "message": message}
     else:
-        message = f"no running agent or preprocess tasks found for conversation_id {conversation_id}"
+        message = f"no running agent or preprocess tasks found for user_id {user_id}, conversation_id {conversation_id}"
         logging.error(message)
         return {"status": "error", "message": message}
 

backend/services/northbound_service.py

@@ -217,7 +217,7 @@ async def stop_chat(ctx: NorthboundContext, external_conversation_id: str) -> Di
     try:
         internal_id = await to_internal_conversation_id(external_conversation_id)
 
-        stop_result = stop_agent_tasks(internal_id)
+        stop_result = stop_agent_tasks(internal_id, ctx.user_id)
         return {"message": stop_result.get("message", "success"), "data": external_conversation_id, "requestId": ctx.request_id}
     except Exception as e:
         raise Exception(f"Failed to stop chat for external conversation id {external_conversation_id}: {str(e)}")