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

Author: ljy65535

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
 
 
@@ -881,13 +881,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(
@@ -900,11 +900,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

@@ -208,7 +208,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)}")

test/backend/app/test_agent_app.py

@@ -105,26 +105,42 @@ async def mock_stream():
 
 def test_agent_stop_api_success(mocker, mock_conversation_id):
     """Test agent_stop_api success case."""
+    # Mock the authentication function to return user_id
+    mock_get_user_id = mocker.patch("apps.agent_app.get_current_user_id")
+    mock_get_user_id.return_value = ("test_user_id", "test_tenant_id")
+    
     mock_stop_tasks = mocker.patch("apps.agent_app.stop_agent_tasks")
     mock_stop_tasks.return_value = {"status": "success"}
 
-    response = client.get(f"/agent/stop/{mock_conversation_id}")
+    response = client.get(
+        f"/agent/stop/{mock_conversation_id}",
+        headers={"Authorization": "Bearer test_token"}
+    )
 
     assert response.status_code == 200
-    mock_stop_tasks.assert_called_once_with(mock_conversation_id)
+    mock_get_user_id.assert_called_once_with("Bearer test_token")
+    mock_stop_tasks.assert_called_once_with(mock_conversation_id, "test_user_id")
     assert response.json()["status"] == "success"
 
 
 def test_agent_stop_api_not_found(mocker, mock_conversation_id):
     """Test agent_stop_api not found case."""
+    # Mock the authentication function to return user_id
+    mock_get_user_id = mocker.patch("apps.agent_app.get_current_user_id")
+    mock_get_user_id.return_value = ("test_user_id", "test_tenant_id")
+    
     mock_stop_tasks = mocker.patch("apps.agent_app.stop_agent_tasks")
     mock_stop_tasks.return_value = {"status": "error"}  # Simulate not found
 
-    response = client.get(f"/agent/stop/{mock_conversation_id}")
+    response = client.get(
+        f"/agent/stop/{mock_conversation_id}",
+        headers={"Authorization": "Bearer test_token"}
+    )
 
     # The app should raise HTTPException for non-success status
     assert response.status_code == 400
-    mock_stop_tasks.assert_called_once_with(mock_conversation_id)
+    mock_get_user_id.assert_called_once_with("Bearer test_token")
+    mock_stop_tasks.assert_called_once_with(mock_conversation_id, "test_user_id")
     assert "no running agent or preprocess tasks found" in response.json()[
         "detail"]
 

test/backend/services/test_agent_service.py

@@ -1838,7 +1838,7 @@ async def test_prepare_agent_run(
         "test_user", "test_tenant", 1)
     mock_create_run_info.assert_called_once()
     mock_agent_run_manager.register_agent_run.assert_called_once_with(
-        123, mock_run_info)
+        123, mock_run_info, "test_user")
 
 
 @patch('backend.services.agent_service.submit')
@@ -1933,21 +1933,24 @@ def test_stop_agent_tasks(mock_preprocess_manager, mock_agent_run_manager):
     # Test both stopped
     mock_agent_run_manager.stop_agent_run.return_value = True
     mock_preprocess_manager.stop_preprocess_tasks.return_value = True
-    result = stop_agent_tasks(123)
+
+    result = stop_agent_tasks(123, "test_user")
     assert result["status"] == "success"
     assert "successfully stopped agent run and preprocess tasks" in result["message"]
 
+    mock_agent_run_manager.stop_agent_run.assert_called_once_with(123, "test_user")
+
     # Test only agent stopped
     mock_agent_run_manager.stop_agent_run.return_value = True
     mock_preprocess_manager.stop_preprocess_tasks.return_value = False
-    result = stop_agent_tasks(123)
+    result = stop_agent_tasks(123, "test_user")
     assert result["status"] == "success"
     assert "successfully stopped agent run" in result["message"]
 
     # Test neither stopped
     mock_agent_run_manager.stop_agent_run.return_value = False
     mock_preprocess_manager.stop_preprocess_tasks.return_value = False
-    result = stop_agent_tasks(123)
+    result = stop_agent_tasks(123, "test_user")
     assert result["status"] == "error"
     assert "no running agent or preprocess tasks found" in result["message"]
 
@@ -2340,11 +2343,11 @@ def fake_save_messages(*args, **kwargs):
         raising=False,
     )
 
-    # Mock unregister
     unregister_called = {}
 
-    def fake_unregister(conv_id):
+    def fake_unregister(conv_id, user_id):
         unregister_called["conv_id"] = conv_id
+        unregister_called["user_id"] = user_id
 
     monkeypatch.setattr(
         "backend.services.agent_service.agent_run_manager.unregister_agent_run",
@@ -2365,6 +2368,7 @@ def fake_unregister(conv_id):
     ]  # Prefix added in helper
     assert save_calls, "save_messages should have been called for assistant messages"
     assert unregister_called.get("conv_id") == 999
+    assert unregister_called.get("user_id") == "u"
 
 
 @pytest.mark.asyncio
@@ -2386,10 +2390,11 @@ async def failing_agent_run(*_, **__):
         "backend.services.agent_service.agent_run", failing_agent_run, raising=False
     )
 
-    called = {"unregistered": None}
+    called = {"unregistered": None, "user_id": None}
 
-    def fake_unregister(conv_id):
+    def fake_unregister(conv_id, user_id):
         called["unregistered"] = conv_id
+        called["user_id"] = user_id
 
     monkeypatch.setattr(
         "backend.services.agent_service.agent_run_manager.unregister_agent_run",
@@ -2408,6 +2413,7 @@ def fake_unregister(conv_id):
     assert collected and collected[0].startswith(
         "data: {") and "\"type\": \"error\"" in collected[0]
     assert called["unregistered"] == 1001
+    assert called["user_id"] == "u"
 
 
 @pytest.mark.asyncio
@@ -2692,13 +2698,13 @@ async def test_generate_stream_no_memory_registers_and_streams(monkeypatch):
         AsyncMock(return_value=MagicMock()),
         raising=False,
     )
-
-    # Capture register
+    
     registered = {}
 
-    def fake_register(conv_id, run_info):
+    def fake_register(conv_id, run_info, user_id):
         registered["conv_id"] = conv_id
         registered["run_info"] = run_info
+        registered["user_id"] = user_id
 
     monkeypatch.setattr(
         "backend.services.agent_service.agent_run_manager.register_agent_run",
@@ -2725,6 +2731,8 @@ async def fake_stream_chunks(*_, **__):
         collected.append(d)
 
     assert registered.get("conv_id") == 555
+    assert registered.get("user_id") == "u"
+    assert registered.get("run_info") is not None
     assert collected == ["data: body1\n\n", "data: body2\n\n"]
 
 

test/backend/services/test_northbound_service.py

@@ -263,6 +263,8 @@ async def test_stop_chat_success(ctx):
     assert result["data"] == "ext-777"
     assert result["requestId"] == "req-1"
 
+    agent_service_mod.stop_agent_tasks.assert_called_once_with(777, "user-1")
+
 
 @pytest.mark.asyncio
 async def test_list_conversations_maps_ids(ctx):