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:
2026-01-03 17:55:34 +01:00
parent 520c06175e
commit caf283bed2
9 changed files with 1782 additions and 92 deletions

View File

@@ -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,