Skip to content

Commit c786e19

Browse files
committed
limit agent creation based on tiers
1 parent f1f8f82 commit c786e19

File tree

18 files changed

+450
-42
lines changed

18 files changed

+450
-42
lines changed

backend/agent/api.py

Lines changed: 47 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1767,7 +1767,6 @@ async def import_agent_from_json(
17671767
request: JsonImportRequestModel,
17681768
user_id: str = Depends(get_current_user_id_from_jwt)
17691769
):
1770-
"""Import an agent from JSON with credential mappings"""
17711770
logger.info(f"Importing agent from JSON - user: {user_id}")
17721771

17731772
if not await is_enabled("custom_agents"):
@@ -1776,6 +1775,21 @@ async def import_agent_from_json(
17761775
detail="Custom agents currently disabled. This feature is not available at the moment."
17771776
)
17781777

1778+
client = await db.client
1779+
from .utils import check_agent_count_limit
1780+
limit_check = await check_agent_count_limit(client, user_id)
1781+
1782+
if not limit_check['can_create']:
1783+
error_detail = {
1784+
"message": f"Maximum of {limit_check['limit']} agents allowed for your current plan. You have {limit_check['current_count']} agents.",
1785+
"current_count": limit_check['current_count'],
1786+
"limit": limit_check['limit'],
1787+
"tier_name": limit_check['tier_name'],
1788+
"error_code": "AGENT_LIMIT_EXCEEDED"
1789+
}
1790+
logger.warning(f"Agent limit exceeded for account {user_id}: {limit_check['current_count']}/{limit_check['limit']} agents")
1791+
raise HTTPException(status_code=402, detail=error_detail)
1792+
17791793
try:
17801794
from agent.json_import_service import JsonImportService, JsonImportRequest
17811795
import_service = JsonImportService(db)
@@ -1817,6 +1831,20 @@ async def create_agent(
18171831
)
18181832
client = await db.client
18191833

1834+
from .utils import check_agent_count_limit
1835+
limit_check = await check_agent_count_limit(client, user_id)
1836+
1837+
if not limit_check['can_create']:
1838+
error_detail = {
1839+
"message": f"Maximum of {limit_check['limit']} agents allowed for your current plan. You have {limit_check['current_count']} agents.",
1840+
"current_count": limit_check['current_count'],
1841+
"limit": limit_check['limit'],
1842+
"tier_name": limit_check['tier_name'],
1843+
"error_code": "AGENT_LIMIT_EXCEEDED"
1844+
}
1845+
logger.warning(f"Agent limit exceeded for account {user_id}: {limit_check['current_count']}/{limit_check['limit']} agents")
1846+
raise HTTPException(status_code=402, detail=error_detail)
1847+
18201848
try:
18211849
if agent_data.is_default:
18221850
await client.table('agents').update({"is_default": False}).eq("account_id", user_id).eq("is_default", True).execute()
@@ -1874,6 +1902,10 @@ async def create_agent(
18741902
await client.table('agents').delete().eq('agent_id', agent['agent_id']).execute()
18751903
raise HTTPException(status_code=500, detail="Failed to create initial version")
18761904

1905+
# Invalidate agent count cache after successful creation
1906+
from utils.cache import Cache
1907+
await Cache.invalidate(f"agent_count_limit:{user_id}")
1908+
18771909
logger.info(f"Created agent {agent['agent_id']} with v1 for user: {user_id}")
18781910
return AgentResponse(
18791911
agent_id=agent['agent_id'],
@@ -2275,7 +2307,20 @@ async def delete_agent(agent_id: str, user_id: str = Depends(get_current_user_id
22752307
if agent['is_default']:
22762308
raise HTTPException(status_code=400, detail="Cannot delete default agent")
22772309

2278-
await client.table('agents').delete().eq('agent_id', agent_id).execute()
2310+
if agent.get('metadata', {}).get('is_suna_default', False):
2311+
raise HTTPException(status_code=400, detail="Cannot delete Suna default agent")
2312+
2313+
delete_result = await client.table('agents').delete().eq('agent_id', agent_id).execute()
2314+
2315+
if not delete_result.data:
2316+
logger.warning(f"No agent was deleted for agent_id: {agent_id}, user_id: {user_id}")
2317+
raise HTTPException(status_code=403, detail="Unable to delete agent - permission denied or agent not found")
2318+
2319+
try:
2320+
from utils.cache import Cache
2321+
await Cache.invalidate(f"agent_count_limit:{user_id}")
2322+
except Exception as cache_error:
2323+
logger.warning(f"Cache invalidation failed for user {user_id}: {str(cache_error)}")
22792324

22802325
logger.info(f"Successfully deleted agent: {agent_id}")
22812326
return {"message": "Agent deleted successfully"}

backend/agent/json_import_service.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -114,6 +114,9 @@ async def import_json(self, request: JsonImportRequest) -> JsonImportResult:
114114
request.custom_system_prompt or json_data.get('system_prompt', '')
115115
)
116116

117+
from utils.cache import Cache
118+
await Cache.invalidate(f"agent_count_limit:{request.account_id}")
119+
117120
logger.info(f"Successfully imported agent {agent_id} from JSON")
118121

119122
return JsonImportResult(

backend/agent/utils.py

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -138,3 +138,63 @@ async def check_agent_run_limit(client, account_id: str) -> Dict[str, Any]:
138138
'running_count': 0,
139139
'running_thread_ids': []
140140
}
141+
142+
143+
async def check_agent_count_limit(client, account_id: str) -> Dict[str, Any]:
144+
try:
145+
try:
146+
result = await Cache.get(f"agent_count_limit:{account_id}")
147+
if result:
148+
logger.debug(f"Cache hit for agent count limit: {account_id}")
149+
return result
150+
except Exception as cache_error:
151+
logger.warning(f"Cache read failed for agent count limit {account_id}: {str(cache_error)}")
152+
153+
agents_result = await client.table('agents').select('agent_id, metadata').eq('account_id', account_id).execute()
154+
155+
non_suna_agents = []
156+
for agent in agents_result.data or []:
157+
metadata = agent.get('metadata', {}) or {}
158+
is_suna_default = metadata.get('is_suna_default', False)
159+
if not is_suna_default:
160+
non_suna_agents.append(agent)
161+
162+
current_count = len(non_suna_agents)
163+
logger.debug(f"Account {account_id} has {current_count} custom agents (excluding Suna defaults)")
164+
165+
try:
166+
from services.billing import get_subscription_tier
167+
tier_name = await get_subscription_tier(client, account_id)
168+
logger.debug(f"Account {account_id} subscription tier: {tier_name}")
169+
except Exception as billing_error:
170+
logger.warning(f"Could not get subscription tier for {account_id}: {str(billing_error)}, defaulting to free")
171+
tier_name = 'free'
172+
173+
agent_limit = config.AGENT_LIMITS.get(tier_name, config.AGENT_LIMITS['free'])
174+
175+
can_create = current_count < agent_limit
176+
177+
result = {
178+
'can_create': can_create,
179+
'current_count': current_count,
180+
'limit': agent_limit,
181+
'tier_name': tier_name
182+
}
183+
184+
try:
185+
await Cache.set(f"agent_count_limit:{account_id}", result, ttl=300)
186+
except Exception as cache_error:
187+
logger.warning(f"Cache write failed for agent count limit {account_id}: {str(cache_error)}")
188+
189+
logger.info(f"Account {account_id} has {current_count}/{agent_limit} agents (tier: {tier_name}) - can_create: {can_create}")
190+
191+
return result
192+
193+
except Exception as e:
194+
logger.error(f"Error checking agent count limit for account {account_id}: {str(e)}", exc_info=True)
195+
return {
196+
'can_create': True,
197+
'current_count': 0,
198+
'limit': config.AGENT_LIMITS['free'],
199+
'tier_name': 'free'
200+
}

backend/services/billing.py

Lines changed: 24 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -592,6 +592,30 @@ async def can_use_model(client, user_id: str, model_name: str):
592592

593593
return False, f"Your current subscription plan does not include access to {model_name}. Please upgrade your subscription or choose from your available models: {', '.join(allowed_models)}", allowed_models
594594

595+
async def get_subscription_tier(client, user_id: str) -> str:
596+
try:
597+
subscription = await get_user_subscription(user_id)
598+
599+
if not subscription:
600+
return 'free'
601+
602+
price_id = None
603+
if subscription.get('items') and subscription['items'].get('data') and len(subscription['items']['data']) > 0:
604+
price_id = subscription['items']['data'][0]['price']['id']
605+
else:
606+
price_id = subscription.get('price_id', config.STRIPE_FREE_TIER_ID)
607+
608+
tier_info = SUBSCRIPTION_TIERS.get(price_id)
609+
if tier_info:
610+
return tier_info['name']
611+
612+
logger.warning(f"Unknown price_id {price_id} for user {user_id}, defaulting to free tier")
613+
return 'free'
614+
615+
except Exception as e:
616+
logger.error(f"Error getting subscription tier for user {user_id}: {str(e)}")
617+
return 'free'
618+
595619
async def check_billing_status(client, user_id: str) -> Tuple[bool, str, Optional[Dict]]:
596620
"""
597621
Check if a user can run agents based on their subscription and usage.
@@ -634,8 +658,6 @@ async def check_billing_status(client, user_id: str) -> Tuple[bool, str, Optiona
634658
# Calculate current month's usage
635659
current_usage = await calculate_monthly_usage(client, user_id)
636660

637-
# TODO: also do user's AAL check
638-
# Check if within limits
639661
if current_usage >= tier_info['cost']:
640662
return False, f"Monthly limit of {tier_info['cost']} dollars reached. Please upgrade your plan or wait until next month.", subscription
641663

backend/templates/api.py

Lines changed: 15 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -325,15 +325,22 @@ async def install_template(
325325
request: InstallTemplateRequest,
326326
user_id: str = Depends(get_current_user_id_from_jwt)
327327
):
328-
"""
329-
Install a template as a new agent instance.
330-
331-
Requires:
332-
- User must have access to the template (own it or it's public)
333-
"""
334328
try:
335-
# Validate template access first
336-
template = await validate_template_access_and_get(request.template_id, user_id)
329+
await validate_template_access_and_get(request.template_id, user_id)
330+
client = await db.client
331+
from agent.utils import check_agent_count_limit
332+
limit_check = await check_agent_count_limit(client, user_id)
333+
334+
if not limit_check['can_create']:
335+
error_detail = {
336+
"message": f"Maximum of {limit_check['limit']} agents allowed for your current plan. You have {limit_check['current_count']} agents.",
337+
"current_count": limit_check['current_count'],
338+
"limit": limit_check['limit'],
339+
"tier_name": limit_check['tier_name'],
340+
"error_code": "AGENT_LIMIT_EXCEEDED"
341+
}
342+
logger.warning(f"Agent limit exceeded for account {user_id}: {limit_check['current_count']}/{limit_check['limit']} agents")
343+
raise HTTPException(status_code=402, detail=error_detail)
337344

338345
logger.info(f"User {user_id} installing template {request.template_id}")
339346

@@ -362,7 +369,6 @@ async def install_template(
362369
)
363370

364371
except HTTPException:
365-
# Re-raise HTTP exceptions from our validation functions
366372
raise
367373
except TemplateInstallationError as e:
368374
logger.warning(f"Template installation failed: {e}")

backend/templates/installation_service.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,9 @@ async def install_template(self, request: TemplateInstallationRequest) -> Templa
107107

108108
await self._increment_download_count(template.template_id)
109109

110+
from utils.cache import Cache
111+
await Cache.invalidate(f"agent_count_limit:{request.account_id}")
112+
110113
agent_name = request.instance_name or f"{template.name} (from marketplace)"
111114
logger.info(f"Successfully installed template {template.template_id} as agent {agent_id}")
112115

backend/utils/config.py

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -270,6 +270,30 @@ def STRIPE_TIER_25_170_YEARLY_COMMITMENT_ID(self) -> str:
270270

271271
# Agent execution limits (can be overridden via environment variable)
272272
_MAX_PARALLEL_AGENT_RUNS_ENV: Optional[str] = None
273+
274+
# Agent limits per billing tier
275+
AGENT_LIMITS = {
276+
'free': 2,
277+
'tier_2_20': 5,
278+
'tier_6_50': 20,
279+
'tier_12_100': 20,
280+
'tier_25_200': 100,
281+
'tier_50_400': 100,
282+
'tier_125_800': 100,
283+
'tier_200_1000': 100,
284+
# Yearly plans have same limits as monthly
285+
'tier_2_20_yearly': 5,
286+
'tier_6_50_yearly': 20,
287+
'tier_12_100_yearly': 20,
288+
'tier_25_200_yearly': 100,
289+
'tier_50_400_yearly': 100,
290+
'tier_125_800_yearly': 100,
291+
'tier_200_1000_yearly': 100,
292+
# Yearly commitment plans
293+
'tier_2_17_yearly_commitment': 5,
294+
'tier_6_42_yearly_commitment': 20,
295+
'tier_25_170_yearly_commitment': 100,
296+
}
273297

274298
@property
275299
def MAX_PARALLEL_AGENT_RUNS(self) -> int:

frontend/src/app/(dashboard)/agents/page.tsx

Lines changed: 20 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,8 @@ import { PublishDialog } from '@/components/agents/custom-agents-page/publish-di
2323
import { LoadingSkeleton } from '@/components/agents/custom-agents-page/loading-skeleton';
2424
import { NewAgentDialog } from '@/components/agents/new-agent-dialog';
2525
import { MarketplaceAgentPreviewDialog } from '@/components/agents/marketplace-agent-preview-dialog';
26+
import { AgentCountLimitDialog } from '@/components/agents/agent-count-limit-dialog';
27+
import { AgentCountLimitError } from '@/lib/api';
2628

2729
type ViewMode = 'grid' | 'list';
2830
type AgentSortOption = 'name' | 'created_at' | 'updated_at' | 'tools_count';
@@ -86,6 +88,8 @@ export default function AgentsPage() {
8688

8789
const [publishingAgentId, setPublishingAgentId] = useState<string | null>(null);
8890
const [showNewAgentDialog, setShowNewAgentDialog] = useState(false);
91+
const [showAgentLimitDialog, setShowAgentLimitDialog] = useState(false);
92+
const [agentLimitError, setAgentLimitError] = useState<AgentCountLimitError | null>(null);
8993

9094
const activeTab = useMemo(() => {
9195
return searchParams.get('tab') || 'my-agents';
@@ -146,7 +150,6 @@ export default function AgentsPage() {
146150

147151
if (marketplaceTemplates) {
148152
marketplaceTemplates.forEach(template => {
149-
150153
const item: MarketplaceTemplate = {
151154
id: template.template_id,
152155
creator_id: template.creator_id,
@@ -174,14 +177,12 @@ export default function AgentsPage() {
174177
item.creator_name?.toLowerCase().includes(searchLower);
175178
})();
176179

177-
if (!matchesSearch) return; // Skip items that don't match search
180+
if (!matchesSearch) return;
178181

179-
// Always add user's own templates to mineItems for the "mine" filter
180182
if (user?.id === template.creator_id) {
181183
mineItems.push(item);
182184
}
183185

184-
// Categorize all templates (including user's own) for the "all" view
185186
if (template.is_kortix_team) {
186187
kortixItems.push(item);
187188
} else {
@@ -392,6 +393,12 @@ export default function AgentsPage() {
392393
} catch (error: any) {
393394
console.error('Installation error:', error);
394395

396+
if (error instanceof AgentCountLimitError) {
397+
setAgentLimitError(error);
398+
setShowAgentLimitDialog(true);
399+
return;
400+
}
401+
395402
if (error.message?.includes('already in your library')) {
396403
toast.error('This agent is already in your library');
397404
} else if (error.message?.includes('Credential profile not found')) {
@@ -628,6 +635,15 @@ export default function AgentsPage() {
628635
onInstall={handlePreviewInstall}
629636
isInstalling={installingItemId === selectedItem?.id}
630637
/>
638+
{agentLimitError && (
639+
<AgentCountLimitDialog
640+
open={showAgentLimitDialog}
641+
onOpenChange={setShowAgentLimitDialog}
642+
currentCount={agentLimitError.detail.current_count}
643+
limit={agentLimitError.detail.limit}
644+
tierName={agentLimitError.detail.tier_name}
645+
/>
646+
)}
631647
</div>
632648
</div>
633649
);

0 commit comments

Comments
 (0)