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:
@@ -20,13 +20,13 @@ from app.services.mcp.routing import ToolInfo, ToolResult
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def reset_registry():
|
||||
async def reset_registry():
|
||||
"""Reset the singleton registry before and after each test."""
|
||||
MCPServerRegistry.reset_instance()
|
||||
reset_mcp_client()
|
||||
await reset_mcp_client()
|
||||
yield
|
||||
MCPServerRegistry.reset_instance()
|
||||
reset_mcp_client()
|
||||
await reset_mcp_client()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
@@ -388,7 +388,8 @@ class TestModuleLevelFunctions:
|
||||
mock_shutdown.return_value = None
|
||||
await shutdown_mcp_client()
|
||||
|
||||
def test_reset_mcp_client(self, reset_registry):
|
||||
@pytest.mark.asyncio
|
||||
async def test_reset_mcp_client(self, reset_registry):
|
||||
"""Test resetting the global client."""
|
||||
reset_mcp_client()
|
||||
await reset_mcp_client()
|
||||
# Should not raise
|
||||
|
||||
Reference in New Issue
Block a user