♻️ Refresh the tool list when backend service is initialized #1250

Author: Phinease

backend/database/user_tenant_db.py

@@ -3,6 +3,7 @@
 """
 from typing import Any, Dict, Optional
 
+from consts.const import DEFAULT_TENANT_ID
 from database.client import as_dict, get_db_session
 from database.db_models import UserTenant
 
@@ -28,6 +29,27 @@ def get_user_tenant_by_user_id(user_id: str) -> Optional[Dict[str, Any]]:
         return None
 
 
+def get_all_tenant_ids() -> list[str]:
+    """
+    Get all unique tenant IDs from the database
+    
+    Returns:
+        list[str]: List of unique tenant IDs
+    """
+    with get_db_session() as session:
+        result = session.query(UserTenant.tenant_id).filter(
+            UserTenant.delete_flag == "N"
+        ).distinct().all()
+        
+        tenant_ids = [row[0] for row in result]
+        
+        # Add default tenant_id if not already in the list
+        if DEFAULT_TENANT_ID not in tenant_ids:
+            tenant_ids.append(DEFAULT_TENANT_ID)
+        
+        return tenant_ids
+
+
 def insert_user_tenant(user_id: str, tenant_id: str):
     """
     Insert user tenant relationship

backend/main_service.py

@@ -1,19 +1,39 @@
 import uvicorn
 import logging
 import warnings
+import asyncio
 warnings.filterwarnings("ignore", category=UserWarning)
 
 from dotenv import load_dotenv
 load_dotenv()
 
 from apps.base_app import app
 from utils.logging_utils import configure_logging, configure_elasticsearch_logging
+from services.tool_configuration_service import initialize_tools_on_startup
 
 
 configure_logging(logging.INFO)
 configure_elasticsearch_logging()
 logger = logging.getLogger("main_service")
 
+
+async def startup_initialization():
+    """
+    Perform initialization tasks during server startup
+    """
+    logger.info("Starting server initialization...")
+    
+    try:
+        # Initialize tools on startup - service layer handles detailed logging
+        await initialize_tools_on_startup()
+        logger.info("Server initialization completed successfully!")
+            
+    except Exception as e:
+        logger.error(f"Server initialization failed: {str(e)}")
+        # Don't raise the exception to allow server to start even if initialization fails
+        logger.warning("Server will continue to start despite initialization issues")
+
+
 if __name__ == "__main__":
-    # Scan agents and update to database
+    asyncio.run(startup_initialization())
     uvicorn.run(app, host="0.0.0.0", port=5010, log_level="info")

backend/services/tool_configuration_service.py

@@ -1,3 +1,4 @@
+import asyncio
 import importlib
 import inspect
 import json
@@ -10,7 +11,7 @@
 import jsonref
 from mcpadapt.smolagents_adapter import _sanitize_function_name
 
-from consts.const import LOCAL_MCP_SERVER
+from consts.const import DEFAULT_USER_ID, LOCAL_MCP_SERVER
 from consts.exceptions import MCPConnectionError
 from consts.model import ToolInstanceInfoRequest, ToolInfo, ToolSourceEnum
 from database.remote_mcp_db import get_mcp_records_by_tenant
@@ -20,6 +21,7 @@
     query_tool_instances_by_id,
     update_tool_table_from_scan_tool_list
 )
+from database.user_tenant_db import get_all_tenant_ids
 
 logger = logging.getLogger("tool_configuration_service")
 
@@ -346,4 +348,69 @@ async def list_all_tools(tenant_id: str):
         }
         formatted_tools.append(formatted_tool)
 
-    return formatted_tools
+    return formatted_tools
+
+
+async def initialize_tools_on_startup():
+    """
+    Initialize and scan all tools during server startup for all tenants
+    
+    This function scans all available tools (local, LangChain, and MCP) 
+    and updates the database with the latest tool information for all tenants.
+    """
+    
+    logger.info("Starting tool initialization on server startup...")
+    
+    try:
+        # Get all tenant IDs from the database
+        tenant_ids = get_all_tenant_ids()
+        
+        if not tenant_ids:
+            logger.warning("No tenants found in database, skipping tool initialization")
+            return
+        
+        logger.info(f"Found {len(tenant_ids)} tenants: {tenant_ids}")
+        
+        total_tools = 0
+        successful_tenants = 0
+        failed_tenants = []
+        
+        # Process each tenant
+        for tenant_id in tenant_ids:
+            try:
+                logger.info(f"Initializing tools for tenant: {tenant_id}")
+                
+                # Add timeout to prevent hanging during startup
+                try:
+                    await asyncio.wait_for(
+                        update_tool_list(tenant_id=tenant_id, user_id=DEFAULT_USER_ID),
+                        timeout=60.0  # 60 seconds timeout per tenant
+                    )
+                    
+                    # Get the count of tools for this tenant
+                    tools_info = query_all_tools(tenant_id)
+                    tenant_tool_count = len(tools_info)
+                    total_tools += tenant_tool_count
+                    successful_tenants += 1
+                    
+                    logger.info(f"Tenant {tenant_id}: {tenant_tool_count} tools initialized")
+                    
+                except asyncio.TimeoutError:
+                    logger.error(f"Tool initialization timed out for tenant {tenant_id}")
+                    failed_tenants.append(f"{tenant_id} (timeout)")
+                    
+            except Exception as e:
+                logger.error(f"Tool initialization failed for tenant {tenant_id}: {str(e)}")
+                failed_tenants.append(f"{tenant_id} (error: {str(e)})")
+        
+        # Log final results
+        logger.info(f"Tool initialization completed!")
+        logger.info(f"Total tools available across all tenants: {total_tools}")
+        logger.info(f"Successfully processed: {successful_tenants}/{len(tenant_ids)} tenants")
+        
+        if failed_tenants:
+            logger.warning(f"Failed tenants: {', '.join(failed_tenants)}")
+        
+    except Exception as e:
+        logger.error(f"❌ Tool initialization failed: {str(e)}")
+        raise

test/backend/database/test_user_tenant_db.py

@@ -66,7 +66,8 @@
 # Now import the functions to be tested
 from backend.database.user_tenant_db import (
     get_user_tenant_by_user_id,
-    insert_user_tenant
+    insert_user_tenant,
+    get_all_tenant_ids
 )
 
 class MockUserTenant:
@@ -301,4 +302,50 @@ def test_get_user_tenant_by_user_id_with_deleted_record(monkeypatch, mock_sessio
 
     assert result is None
     # Verify that the filter was called with correct conditions
-    query.filter.assert_called_once() 
+    query.filter.assert_called_once()
+
+
+def test_get_all_tenant_ids_empty_database(monkeypatch, mock_session):
+    """Test get_all_tenant_ids when database is empty - should return only DEFAULT_TENANT_ID"""
+    session, query = mock_session
+    
+    # Mock empty database result
+    query.filter.return_value.distinct.return_value.all.return_value = []
+    
+    mock_ctx = MagicMock()
+    mock_ctx.__enter__.return_value = session
+    mock_ctx.__exit__.return_value = None
+    monkeypatch.setattr("backend.database.user_tenant_db.get_db_session", lambda: mock_ctx)
+    
+    result = get_all_tenant_ids()
+    
+    assert result == ["default_tenant"]  # DEFAULT_TENANT_ID from consts_mock
+    assert len(result) == 1
+
+
+def test_get_all_tenant_ids_with_existing_tenants(monkeypatch, mock_session):
+    """Test get_all_tenant_ids with existing tenants - should include all plus DEFAULT_TENANT_ID"""
+    session, query = mock_session
+    
+    # Mock database result with existing tenants
+    mock_tenants = [
+        ("tenant_1",),
+        ("tenant_2",),
+        ("tenant_3",)
+    ]
+    query.filter.return_value.distinct.return_value.all.return_value = mock_tenants
+    
+    mock_ctx = MagicMock()
+    mock_ctx.__enter__.return_value = session
+    mock_ctx.__exit__.return_value = None
+    monkeypatch.setattr("backend.database.user_tenant_db.get_db_session", lambda: mock_ctx)
+    
+    result = get_all_tenant_ids()
+    
+    assert len(result) == 4  # 3 existing + 1 default
+    assert "tenant_1" in result
+    assert "tenant_2" in result
+    assert "tenant_3" in result
+    assert "default_tenant" in result  # DEFAULT_TENANT_ID from consts_mock
+    # Should not duplicate DEFAULT_TENANT_ID
+    assert result.count("default_tenant") == 1