Merge pull request #1267 from escapade-mckv/main

UI fixes

Author: escapade-mckv

backend/agent/api.py

@@ -1790,7 +1790,6 @@ async def import_agent_from_json(
     request: JsonImportRequestModel,
     user_id: str = Depends(get_current_user_id_from_jwt)
 ):
-    """Import an agent from JSON with credential mappings"""
     logger.info(f"Importing agent from JSON - user: {user_id}")
 
     if not await is_enabled("custom_agents"):
@@ -1799,6 +1798,21 @@ async def import_agent_from_json(
             detail="Custom agents currently disabled. This feature is not available at the moment."
         )
 
+    client = await db.client
+    from .utils import check_agent_count_limit
+    limit_check = await check_agent_count_limit(client, user_id)
+    
+    if not limit_check['can_create']:
+        error_detail = {
+            "message": f"Maximum of {limit_check['limit']} agents allowed for your current plan. You have {limit_check['current_count']} agents.",
+            "current_count": limit_check['current_count'],
+            "limit": limit_check['limit'],
+            "tier_name": limit_check['tier_name'],
+            "error_code": "AGENT_LIMIT_EXCEEDED"
+        }
+        logger.warning(f"Agent limit exceeded for account {user_id}: {limit_check['current_count']}/{limit_check['limit']} agents")
+        raise HTTPException(status_code=402, detail=error_detail)
+    
     try:
         from agent.json_import_service import JsonImportService, JsonImportRequest
         import_service = JsonImportService(db)
@@ -1840,6 +1854,20 @@ async def create_agent(
         )
     client = await db.client
 
+    from .utils import check_agent_count_limit
+    limit_check = await check_agent_count_limit(client, user_id)
+    
+    if not limit_check['can_create']:
+        error_detail = {
+            "message": f"Maximum of {limit_check['limit']} agents allowed for your current plan. You have {limit_check['current_count']} agents.",
+            "current_count": limit_check['current_count'],
+            "limit": limit_check['limit'],
+            "tier_name": limit_check['tier_name'],
+            "error_code": "AGENT_LIMIT_EXCEEDED"
+        }
+        logger.warning(f"Agent limit exceeded for account {user_id}: {limit_check['current_count']}/{limit_check['limit']} agents")
+        raise HTTPException(status_code=402, detail=error_detail)
+    
     try:
         if agent_data.is_default:
             await client.table('agents').update({"is_default": False}).eq("account_id", user_id).eq("is_default", True).execute()
@@ -1897,6 +1925,10 @@ async def create_agent(
             await client.table('agents').delete().eq('agent_id', agent['agent_id']).execute()
             raise HTTPException(status_code=500, detail="Failed to create initial version")
 
+        # Invalidate agent count cache after successful creation
+        from utils.cache import Cache
+        await Cache.invalidate(f"agent_count_limit:{user_id}")
+        
         logger.info(f"Created agent {agent['agent_id']} with v1 for user: {user_id}")
         return AgentResponse(
             agent_id=agent['agent_id'],
@@ -2298,7 +2330,20 @@ async def delete_agent(agent_id: str, user_id: str = Depends(get_current_user_id
         if agent['is_default']:
             raise HTTPException(status_code=400, detail="Cannot delete default agent")
 
-        await client.table('agents').delete().eq('agent_id', agent_id).execute()
+        if agent.get('metadata', {}).get('is_suna_default', False):
+            raise HTTPException(status_code=400, detail="Cannot delete Suna default agent")
+        
+        delete_result = await client.table('agents').delete().eq('agent_id', agent_id).execute()
+        
+        if not delete_result.data:
+            logger.warning(f"No agent was deleted for agent_id: {agent_id}, user_id: {user_id}")
+            raise HTTPException(status_code=403, detail="Unable to delete agent - permission denied or agent not found")
+        
+        try:
+            from utils.cache import Cache
+            await Cache.invalidate(f"agent_count_limit:{user_id}")
+        except Exception as cache_error:
+            logger.warning(f"Cache invalidation failed for user {user_id}: {str(cache_error)}")
 
         logger.info(f"Successfully deleted agent: {agent_id}")
         return {"message": "Agent deleted successfully"}

backend/agent/json_import_service.py

@@ -114,6 +114,9 @@ async def import_json(self, request: JsonImportRequest) -> JsonImportResult:
             request.custom_system_prompt or json_data.get('system_prompt', '')
         )
 
+        from utils.cache import Cache
+        await Cache.invalidate(f"agent_count_limit:{request.account_id}")
+        
         logger.info(f"Successfully imported agent {agent_id} from JSON")
 
         return JsonImportResult(

backend/agent/utils.py

@@ -138,3 +138,63 @@ async def check_agent_run_limit(client, account_id: str) -> Dict[str, Any]:
             'running_count': 0,
             'running_thread_ids': []
         }
+
+
+async def check_agent_count_limit(client, account_id: str) -> Dict[str, Any]:
+    try:
+        try:
+            result = await Cache.get(f"agent_count_limit:{account_id}")
+            if result:
+                logger.debug(f"Cache hit for agent count limit: {account_id}")
+                return result
+        except Exception as cache_error:
+            logger.warning(f"Cache read failed for agent count limit {account_id}: {str(cache_error)}")
+
+        agents_result = await client.table('agents').select('agent_id, metadata').eq('account_id', account_id).execute()
+        
+        non_suna_agents = []
+        for agent in agents_result.data or []:
+            metadata = agent.get('metadata', {}) or {}
+            is_suna_default = metadata.get('is_suna_default', False)
+            if not is_suna_default:
+                non_suna_agents.append(agent)
+                
+        current_count = len(non_suna_agents)
+        logger.debug(f"Account {account_id} has {current_count} custom agents (excluding Suna defaults)")
+        
+        try:
+            from services.billing import get_subscription_tier
+            tier_name = await get_subscription_tier(client, account_id)
+            logger.debug(f"Account {account_id} subscription tier: {tier_name}")
+        except Exception as billing_error:
+            logger.warning(f"Could not get subscription tier for {account_id}: {str(billing_error)}, defaulting to free")
+            tier_name = 'free'
+        
+        agent_limit = config.AGENT_LIMITS.get(tier_name, config.AGENT_LIMITS['free'])
+        
+        can_create = current_count < agent_limit
+        
+        result = {
+            'can_create': can_create,
+            'current_count': current_count,
+            'limit': agent_limit,
+            'tier_name': tier_name
+        }
+        
+        try:
+            await Cache.set(f"agent_count_limit:{account_id}", result, ttl=300)
+        except Exception as cache_error:
+            logger.warning(f"Cache write failed for agent count limit {account_id}: {str(cache_error)}")
+        
+        logger.info(f"Account {account_id} has {current_count}/{agent_limit} agents (tier: {tier_name}) - can_create: {can_create}")
+        
+        return result
+        
+    except Exception as e:
+        logger.error(f"Error checking agent count limit for account {account_id}: {str(e)}", exc_info=True)
+        return {
+            'can_create': True,
+            'current_count': 0,
+            'limit': config.AGENT_LIMITS['free'],
+            'tier_name': 'free'
+        }

backend/composio_integration/api.py

@@ -10,7 +10,7 @@
 from .composio_service import (
     get_integration_service,
 )
-from .toolkit_service import ToolkitService
+from .toolkit_service import ToolkitService, ToolsListResponse
 from .composio_profile_service import ComposioProfileService, ComposioProfile
 
 router = APIRouter(prefix="/composio", tags=["composio"])
@@ -21,15 +21,13 @@ def initialize(database: DBConnection):
     global db
     db = database
 
-
 class IntegrateToolkitRequest(BaseModel):
     toolkit_slug: str
     profile_name: Optional[str] = None
     display_name: Optional[str] = None
     mcp_server_name: Optional[str] = None
     save_as_profile: bool = True
 
-
 class IntegrationStatusResponse(BaseModel):
     status: str
     toolkit: str
@@ -40,7 +38,6 @@ class IntegrationStatusResponse(BaseModel):
     profile_id: Optional[str] = None
     redirect_url: Optional[str] = None
 
-
 class CreateProfileRequest(BaseModel):
     toolkit_slug: str
     profile_name: str
@@ -49,6 +46,10 @@ class CreateProfileRequest(BaseModel):
     is_default: bool = False
     initiation_fields: Optional[Dict[str, str]] = None
 
+class ToolsListRequest(BaseModel):
+    toolkit_slug: str
+    limit: int = 50
+    cursor: Optional[str] = None
 
 class ProfileResponse(BaseModel):
     profile_id: str
@@ -406,6 +407,35 @@ async def get_toolkit_icon(
         raise HTTPException(status_code=500, detail="Internal server error")
 
 
+@router.post("/tools/list")
+async def list_toolkit_tools(
+    request: ToolsListRequest,
+    current_user_id: str = Depends(get_current_user_id_from_jwt)
+):
+    try:
+        logger.info(f"User {current_user_id} requesting tools for toolkit: {request.toolkit_slug}")
+        
+        toolkit_service = ToolkitService()
+        tools_response = await toolkit_service.get_toolkit_tools(
+            toolkit_slug=request.toolkit_slug,
+            limit=request.limit,
+            cursor=request.cursor
+        )
+        
+        return {
+            "success": True,
+            "tools": [tool.dict() for tool in tools_response.items],
+            "total_items": tools_response.total_items,
+            "current_page": tools_response.current_page,
+            "total_pages": tools_response.total_pages,
+            "next_cursor": tools_response.next_cursor
+        }
+        
+    except Exception as e:
+        logger.error(f"Failed to list toolkit tools for {request.toolkit_slug}: {e}", exc_info=True)
+        raise HTTPException(status_code=500, detail=f"Failed to get toolkit tools: {str(e)}")
+
+
 @router.get("/health")
 async def health_check() -> Dict[str, str]:
     try:

backend/composio_integration/toolkit_service.py

@@ -48,6 +48,31 @@ class DetailedToolkitInfo(BaseModel):
     base_url: Optional[str] = None
 
 
+class ParameterSchema(BaseModel):
+    properties: Dict[str, Any] = {}
+    required: Optional[List[str]] = None
+
+
+class ToolInfo(BaseModel):
+    slug: str
+    name: str
+    description: str
+    version: str
+    input_parameters: ParameterSchema = ParameterSchema()
+    output_parameters: ParameterSchema = ParameterSchema()
+    scopes: List[str] = []
+    tags: List[str] = []
+    no_auth: bool = False
+
+
+class ToolsListResponse(BaseModel):
+    items: List[ToolInfo]
+    next_cursor: Optional[str] = None
+    total_items: int
+    current_page: int = 1
+    total_pages: int = 1
+
+
 class ToolkitService:
     def __init__(self, api_key: Optional[str] = None):
         self.client = ComposioClient.get_client(api_key)
@@ -388,4 +413,80 @@ async def get_detailed_toolkit_info(self, toolkit_slug: str) -> Optional[Detaile
 
         except Exception as e:
             logger.error(f"Failed to get detailed toolkit info for {toolkit_slug}: {e}", exc_info=True)
-            return None 
+            return None
+
+    async def get_toolkit_tools(self, toolkit_slug: str, limit: int = 50, cursor: Optional[str] = None) -> ToolsListResponse:
+        try:
+            logger.info(f"Fetching tools for toolkit: {toolkit_slug}")
+            
+            params = {
+                "limit": limit,
+                "toolkit_slug": toolkit_slug
+            }
+            
+            if cursor:
+                params["cursor"] = cursor
+            
+            tools_response = self.client.tools.list(**params)
+            
+            if hasattr(tools_response, '__dict__'):
+                response_data = tools_response.__dict__
+            else:
+                response_data = tools_response
+            
+            items = response_data.get('items', [])
+            
+            tools = []
+            for item in items:
+                if hasattr(item, '__dict__'):
+                    tool_data = item.__dict__
+                elif hasattr(item, '_asdict'):
+                    tool_data = item._asdict()
+                else:
+                    tool_data = item
+                
+                input_params_raw = tool_data.get("input_parameters", {})
+                output_params_raw = tool_data.get("output_parameters", {})
+                
+                input_parameters = ParameterSchema()
+                if isinstance(input_params_raw, dict):
+                    input_parameters.properties = input_params_raw.get("properties", input_params_raw)
+                    input_parameters.required = input_params_raw.get("required")
+                
+                output_parameters = ParameterSchema()  
+                if isinstance(output_params_raw, dict):
+                    output_parameters.properties = output_params_raw.get("properties", output_params_raw)
+                    output_parameters.required = output_params_raw.get("required")
+                
+                tool = ToolInfo(
+                    slug=tool_data.get("slug", ""),
+                    name=tool_data.get("name", ""),
+                    description=tool_data.get("description", ""),
+                    version=tool_data.get("version", "1.0.0"),
+                    input_parameters=input_parameters,
+                    output_parameters=output_parameters,
+                    scopes=tool_data.get("scopes", []),
+                    tags=tool_data.get("tags", []),
+                    no_auth=tool_data.get("no_auth", False)
+                )
+                tools.append(tool)
+            
+            result = ToolsListResponse(
+                items=tools,
+                total_items=response_data.get("total_items", len(tools)),
+                total_pages=response_data.get("total_pages", 1),
+                current_page=response_data.get("current_page", 1),
+                next_cursor=response_data.get("next_cursor")
+            )
+            
+            logger.info(f"Successfully fetched {len(tools)} tools for toolkit {toolkit_slug}")
+            return result
+            
+        except Exception as e:
+            logger.error(f"Failed to get tools for toolkit {toolkit_slug}: {e}", exc_info=True)
+            return ToolsListResponse(
+                items=[],
+                total_items=0,
+                current_page=1,
+                total_pages=1
+            )