♻️ refactor tool app/service/db/test case

Author: Phinease

backend/apps/tool_config_app.py

@@ -1,15 +1,16 @@
 import logging
+from http import HTTPStatus
 from typing import Optional
 
 from fastapi import APIRouter, Header, HTTPException
 from fastapi.responses import JSONResponse
 
+from consts.exceptions import MCPConnectionError
 from consts.model import ToolInstanceInfoRequest, ToolInstanceSearchRequest
-from database.tool_db import query_all_tools
 from services.tool_configuration_service import (
     search_tool_info_impl,
     update_tool_info_impl,
-    update_tool_list,
+    update_tool_list, list_all_tools,
 )
 from utils.auth_utils import get_current_user_id
 
@@ -25,21 +26,22 @@ async def list_tools_api(authorization: Optional[str] = Header(None)):
     try:
         _, tenant_id = get_current_user_id(authorization)
         # now only admin can modify the tool, user_id is not used
-        return query_all_tools(tenant_id=tenant_id)
+        return await list_all_tools(tenant_id=tenant_id)
     except Exception as e:
         logging.error(f"Failed to get tool info, error in: {str(e)}")
         raise HTTPException(
-            status_code=500, detail=f"Failed to get tool info, error in: {str(e)}")
+            status_code=HTTPStatus.INTERNAL_SERVER_ERROR, detail=f"Failed to get tool info, error in: {str(e)}")
 
 
 @router.post("/search")
 async def search_tool_info_api(request: ToolInstanceSearchRequest, authorization: Optional[str] = Header(None)):
     try:
-        return search_tool_info_impl(request.agent_id, request.tool_id, authorization)
+        _, tenant_id = get_current_user_id(authorization)
+        return search_tool_info_impl(request.agent_id, request.tool_id, tenant_id)
     except Exception as e:
-        logging.error(f"Failed to update tool, error in: {str(e)}")
+        logging.error(f"Failed to search tool, error in: {str(e)}")
         raise HTTPException(
-            status_code=500, detail=f"Failed to update tool, error in: {str(e)}")
+            status_code=HTTPStatus.INTERNAL_SERVER_ERROR, detail="Failed to search tool info")
 
 
 @router.post("/update")
@@ -48,11 +50,12 @@ async def update_tool_info_api(request: ToolInstanceInfoRequest, authorization:
     Update an existing tool, create or update tool instance
     """
     try:
-        return update_tool_info_impl(request, authorization)
+        user_id, tenant_id = get_current_user_id(authorization)
+        return update_tool_info_impl(request, tenant_id, user_id)
     except Exception as e:
         logging.error(f"Failed to update tool, error in: {str(e)}")
         raise HTTPException(
-            status_code=500, detail=f"Failed to update tool, error in: {str(e)}")
+            status_code=HTTPStatus.INTERNAL_SERVER_ERROR, detail=f"Failed to update tool, error in: {str(e)}")
 
 
 @router.get("/scan_tool")
@@ -64,12 +67,14 @@ async def scan_and_update_tool(
         user_id, tenant_id = get_current_user_id(authorization)
         await update_tool_list(tenant_id=tenant_id, user_id=user_id)
         return JSONResponse(
-            status_code=200,
+            status_code=HTTPStatus.OK,
             content={"message": "Successfully update tool", "status": "success"}
         )
+    except MCPConnectionError as e:
+        logger.error(f"MCP connection failed: {e}")
+        raise HTTPException(
+            status_code=HTTPStatus.SERVICE_UNAVAILABLE, detail="MCP connection failed")
     except Exception as e:
         logger.error(f"Failed to update tool: {e}")
-        return JSONResponse(
-            status_code=400,
-            content={"message": "Failed to update tool", "status": "error"}
-        )
+        raise HTTPException(
+            status_code=HTTPStatus.INTERNAL_SERVER_ERROR, detail="Failed to update tool")

backend/database/tool_db.py

@@ -29,7 +29,6 @@ def create_or_update_tool_by_tool_info(tool_info, tenant_id: str, user_id: str):
     :param user_id: Optional user ID for filtering
     :return: Created or updated ToolInstance object
     """
-
     tool_info_dict = tool_info.__dict__ | {
         "tenant_id": tenant_id, "user_id": user_id}
 
@@ -42,18 +41,14 @@ def create_or_update_tool_by_tool_info(tool_info, tenant_id: str, user_id: str):
             ToolInstance.delete_flag != 'Y',
             ToolInstance.tool_id == tool_info_dict['tool_id']
         )
-
         tool_instance = query.first()
-
         if tool_instance:
             # Update the existing ToolInstance
             for key, value in tool_info_dict.items():
                 if hasattr(tool_instance, key):
                     setattr(tool_instance, key, value)
         else:
             create_tool(tool_info_dict)
-
-        session.flush()
         return tool_instance
 
 
@@ -126,41 +121,37 @@ def update_tool_table_from_scan_tool_list(tenant_id: str, user_id: str, tool_lis
     """
     scan all tools and update the tool table in PG database, remove the duplicate tools
     """
-    try:
-        with get_db_session() as session:
-            # get all existing tools (including complete information)
-            existing_tools = session.query(ToolInfo).filter(ToolInfo.delete_flag != 'Y',
-                                                            ToolInfo.author == tenant_id).all()
-            existing_tool_dict = {
-                f"{tool.name}&{tool.source}": tool for tool in existing_tools}
-            # set all tools to unavailable
-            for tool in existing_tools:
-                tool.is_available = False
-
-            for tool in tool_list:
-                filtered_tool_data = filter_property(tool.__dict__, ToolInfo)
-
-                # check if the tool name is valid
-                is_available = True if re.match(
-                    r'^[a-zA-Z_][a-zA-Z0-9_]*$', tool.name) is not None else False
-
-                if f"{tool.name}&{tool.source}" in existing_tool_dict:
-                    # by tool name and source to update the existing tool
-                    existing_tool = existing_tool_dict[f"{tool.name}&{tool.source}"]
-                    for key, value in filtered_tool_data.items():
-                        setattr(existing_tool, key, value)
-                    existing_tool.updated_by = user_id
-                    existing_tool.is_available = is_available
-                else:
-                    # create new tool
-                    filtered_tool_data.update(
-                        {"created_by": user_id, "updated_by": user_id, "author": tenant_id, "is_available": is_available})
-                    new_tool = ToolInfo(**filtered_tool_data)
-                    session.add(new_tool)
-            session.flush()
-        logger.info("Updated tool table in PG database")
-    except Exception as e:
-        logger.error(f"Updated tool table failed due to {e}")
+    with get_db_session() as session:
+        # get all existing tools (including complete information)
+        existing_tools = session.query(ToolInfo).filter(ToolInfo.delete_flag != 'Y',
+                                                        ToolInfo.author == tenant_id).all()
+        existing_tool_dict = {
+            f"{tool.name}&{tool.source}": tool for tool in existing_tools}
+        # set all tools to unavailable
+        for tool in existing_tools:
+            tool.is_available = False
+
+        for tool in tool_list:
+            filtered_tool_data = filter_property(tool.__dict__, ToolInfo)
+
+            # check if the tool name is valid
+            is_available = True if re.match(
+                r'^[a-zA-Z_][a-zA-Z0-9_]*$', tool.name) is not None else False
+
+            if f"{tool.name}&{tool.source}" in existing_tool_dict:
+                # by tool name and source to update the existing tool
+                existing_tool = existing_tool_dict[f"{tool.name}&{tool.source}"]
+                for key, value in filtered_tool_data.items():
+                    setattr(existing_tool, key, value)
+                existing_tool.updated_by = user_id
+                existing_tool.is_available = is_available
+            else:
+                # create new tool
+                filtered_tool_data.update(
+                    {"created_by": user_id, "updated_by": user_id, "author": tenant_id, "is_available": is_available})
+                new_tool = ToolInfo(**filtered_tool_data)
+                session.add(new_tool)
+    logger.info("Updated tool table in PG database")
 
 
 def add_tool_field(tool_info):
@@ -178,14 +169,14 @@ def add_tool_field(tool_info):
         tool_dict = as_dict(tool)
         tool_dict["params"] = tool_params
 
-        # 合并tool_info和tool_dict
+        # combine tool_info and tool_dict
         tool_info.update(tool_dict)
         return tool_info
 
 
 def search_tools_for_sub_agent(agent_id, tenant_id):
     with get_db_session() as session:
-        # Query if there is an existing ToolInstance
+        # query if there is an existing ToolInstance
         query = session.query(ToolInstance).filter(
             ToolInstance.agent_id == agent_id,
             ToolInstance.tenant_id == tenant_id,
@@ -221,4 +212,3 @@ def delete_tools_by_agent_id(agent_id, tenant_id, user_id):
         ).update({
             ToolInstance.delete_flag: 'Y', 'updated_by': user_id
         })
-        session.commit()

backend/services/tool_configuration_service.py

@@ -11,12 +11,10 @@
 from mcpadapt.smolagents_adapter import _sanitize_function_name
 
 from database.tool_db import create_or_update_tool_by_tool_info, query_tool_instances_by_id, \
-    update_tool_table_from_scan_tool_list
+    update_tool_table_from_scan_tool_list, query_all_tools
 from consts.model import ToolInstanceInfoRequest, ToolInfo, ToolSourceEnum
 from consts.const import LOCAL_MCP_SERVER
 from database.remote_mcp_db import get_mcp_records_by_tenant
-from utils.auth_utils import get_current_user_id
-from fastapi import Header
 
 from consts.exceptions import MCPConnectionError
 
@@ -207,31 +205,23 @@ async def get_all_mcp_tools(tenant_id: str) -> List[ToolInfo]:
     return tools_info
 
 
-def search_tool_info_impl(agent_id: int, tool_id: int, authorization: str = Header(None)):
+def search_tool_info_impl(agent_id: int, tool_id: int, tenant_id: str):
     """
     Search for tool configuration information by agent ID and tool ID
 
     Args:
         agent_id: Agent ID
         tool_id: Tool ID
-        authorization:
+        tenant_id: Tenant ID
 
     Returns:
         Dictionary containing tool parameters and enabled status
 
     Raises:
         ValueError: If database query fails
     """
-    _, tenant_id = get_current_user_id(authorization)
-    try:
-        # now only admin can modify the tool, user_id is not used
-        tool_instance = query_tool_instances_by_id(
-            agent_id, tool_id, tenant_id)
-    except Exception as e:
-        logger.error(
-            f"search_tool_info_impl error in query_tool_instances_by_id, detail: {e}")
-        raise ValueError(
-            f"search_tool_info_impl error in query_tool_instances_by_id, detail: {e}")
+    tool_instance = query_tool_instances_by_id(
+        agent_id, tool_id, tenant_id)
 
     if tool_instance:
         return {
@@ -245,29 +235,21 @@ def search_tool_info_impl(agent_id: int, tool_id: int, authorization: str = Head
         }
 
 
-def update_tool_info_impl(request: ToolInstanceInfoRequest, authorization: str = Header(None)):
+def update_tool_info_impl(tool_info: ToolInstanceInfoRequest, tenant_id: str, user_id: str):
     """
     Update tool configuration information
 
     Args:
-        request: ToolInstanceInfoRequest containing tool configuration data
+        tool_info: ToolInstanceInfoRequest containing tool configuration data
 
     Returns:
         Dictionary containing the updated tool instance
 
     Raises:
         ValueError: If database update fails
     """
-    user_id, tenant_id = get_current_user_id(authorization)
-    try:
-        tool_instance = create_or_update_tool_by_tool_info(
-            request, tenant_id, user_id)
-    except Exception as e:
-        logger.error(
-            f"update_tool_info_impl error in create_or_update_tool, detail: {e}")
-        raise ValueError(
-            f"update_tool_info_impl error in create_or_update_tool, detail: {e}")
-
+    tool_instance = create_or_update_tool_by_tool_info(
+        tool_info, tenant_id, user_id)
     return {
         "tool_instance": tool_instance
     }
@@ -313,7 +295,6 @@ async def get_tool_from_remote_mcp_server(mcp_server_name: str, remote_mcp_serve
             f"failed to get tool from remote MCP server, detail: {e}")
 
 
-
 async def update_tool_list(tenant_id: str, user_id: str):
     """
         Scan and gather all available tools from both local and MCP sources
@@ -326,21 +307,38 @@ async def update_tool_list(tenant_id: str, user_id: str):
             List of ToolInfo objects containing tool metadata
         """
     local_tools = get_local_tools()
-
     # Discover LangChain tools (decorated functions) and include them in the
-    # unified tool list.
     langchain_tools = get_langchain_tools()
 
     try:
         mcp_tools = await get_all_mcp_tools(tenant_id)
     except Exception as e:
         logger.error(f"failed to get all mcp tools, detail: {e}")
-        raise Exception(f"failed to get all mcp tools, detail: {e}")
+        raise MCPConnectionError(f"failed to get all mcp tools, detail: {e}")
 
-    try:
-        update_tool_table_from_scan_tool_list(tenant_id=tenant_id,
-                                              user_id=user_id,
-                                              tool_list=local_tools+mcp_tools+langchain_tools)
-    except Exception as e:
-        logger.error(f"failed to update tool list to PG, detail: {e}")
-        raise Exception(f"failed to update tool list to PG, detail: {e}")
+    update_tool_table_from_scan_tool_list(tenant_id=tenant_id,
+                                          user_id=user_id,
+                                          tool_list=local_tools+mcp_tools+langchain_tools)
+
+
+async def list_all_tools(tenant_id: str):
+    """
+    List all tools for a given tenant
+    """
+    tools_info = query_all_tools(tenant_id)
+    # only return the fields needed
+    formatted_tools = []
+    for tool in tools_info:
+        formatted_tool = {
+            "tool_id": tool.get("tool_id"),
+            "name": tool.get("name"),
+            "description": tool.get("description"),
+            "source": tool.get("source"),
+            "is_available": tool.get("is_available"),
+            "create_time": tool.get("create_time"),
+            "usage": tool.get("usage"),
+            "params": tool.get("params", [])
+        }
+        formatted_tools.append(formatted_tool)
+    
+    return formatted_tools

frontend/services/mcpService.ts

@@ -266,14 +266,8 @@ export const updateToolList = async () => {
       let errorMessage = data.message || t('mcpService.message.updateToolListFailed');
 
       switch (response.status) {
-        case 400:
-          errorMessage = t('mcpService.message.updateToolListBadRequest');
-          break;
-        case 404:
-          errorMessage = t('mcpService.message.resourceNotFound');
-          break;
         case 500:
-          errorMessage = t('mcpService.message.serverInternalError');
+          errorMessage = t('mcpService.message.updateToolListBadRequest');
           break;
         case 503:
           errorMessage = t('mcpService.message.serviceUnavailable');