forked from cardosofelipe/fast-next-template
feat(safety): enhance rate limiting and cost control with alert deduplication and usage tracking
- Added `record_action` in `RateLimiter` for precise tracking of slot consumption post-validation. - Introduced deduplication mechanism for warning alerts in `CostController` to prevent spamming. - Refactored `CostController`'s session and daily budget alert handling for improved clarity. - Implemented test suites for `CostController` and `SafetyGuardian` to validate changes. - Expanded integration testing to cover deduplication, validation, and loop detection edge cases.
This commit is contained in:
@@ -223,7 +223,10 @@ class RateLimiter:
|
||||
action: ActionRequest,
|
||||
) -> tuple[bool, list[RateLimitStatus]]:
|
||||
"""
|
||||
Check all applicable rate limits for an action.
|
||||
Check all applicable rate limits for an action WITHOUT consuming slots.
|
||||
|
||||
Use this during validation to check if action would be allowed.
|
||||
Call record_action() after successful execution to consume slots.
|
||||
|
||||
Args:
|
||||
action: The action to check
|
||||
@@ -235,28 +238,53 @@ class RateLimiter:
|
||||
statuses: list[RateLimitStatus] = []
|
||||
allowed = True
|
||||
|
||||
# Check general actions limit
|
||||
actions_allowed, actions_status = await self.acquire("actions", agent_id)
|
||||
# Check general actions limit (read-only)
|
||||
actions_status = await self.check("actions", agent_id)
|
||||
statuses.append(actions_status)
|
||||
if not actions_allowed:
|
||||
if actions_status.is_limited:
|
||||
allowed = False
|
||||
|
||||
# Check LLM-specific limit for LLM calls
|
||||
if action.action_type.value == "llm_call":
|
||||
llm_allowed, llm_status = await self.acquire("llm_calls", agent_id)
|
||||
llm_status = await self.check("llm_calls", agent_id)
|
||||
statuses.append(llm_status)
|
||||
if not llm_allowed:
|
||||
if llm_status.is_limited:
|
||||
allowed = False
|
||||
|
||||
# Check file ops limit for file operations
|
||||
if action.action_type.value in {"file_read", "file_write", "file_delete"}:
|
||||
file_allowed, file_status = await self.acquire("file_ops", agent_id)
|
||||
file_status = await self.check("file_ops", agent_id)
|
||||
statuses.append(file_status)
|
||||
if not file_allowed:
|
||||
if file_status.is_limited:
|
||||
allowed = False
|
||||
|
||||
return allowed, statuses
|
||||
|
||||
async def record_action(
|
||||
self,
|
||||
action: ActionRequest,
|
||||
) -> None:
|
||||
"""
|
||||
Record an action by consuming rate limit slots.
|
||||
|
||||
Call this AFTER successful execution to properly count the action.
|
||||
|
||||
Args:
|
||||
action: The executed action
|
||||
"""
|
||||
agent_id = action.metadata.agent_id
|
||||
|
||||
# Consume general actions slot
|
||||
await self.acquire("actions", agent_id)
|
||||
|
||||
# Consume LLM-specific slot for LLM calls
|
||||
if action.action_type.value == "llm_call":
|
||||
await self.acquire("llm_calls", agent_id)
|
||||
|
||||
# Consume file ops slot for file operations
|
||||
if action.action_type.value in {"file_read", "file_write", "file_delete"}:
|
||||
await self.acquire("file_ops", agent_id)
|
||||
|
||||
async def require(
|
||||
self,
|
||||
limit_name: str,
|
||||
|
||||
Reference in New Issue
Block a user