mirror of
https://github.com/Significant-Gravitas/AutoGPT.git
synced 2026-04-08 03:00:28 -04:00
Merge branch 'dev' of github.com:Significant-Gravitas/AutoGPT into fix/dry-run-simulation-streaming
This commit is contained in:
@@ -30,9 +30,14 @@ from backend.copilot.model import (
|
||||
from backend.copilot.rate_limit import (
|
||||
CoPilotUsageStatus,
|
||||
RateLimitExceeded,
|
||||
acquire_reset_lock,
|
||||
check_rate_limit,
|
||||
get_daily_reset_count,
|
||||
get_global_rate_limits,
|
||||
get_usage_status,
|
||||
increment_daily_reset_count,
|
||||
release_reset_lock,
|
||||
reset_daily_usage,
|
||||
)
|
||||
from backend.copilot.response_model import StreamError, StreamFinish, StreamHeartbeat
|
||||
from backend.copilot.tools.e2b_sandbox import kill_sandbox
|
||||
@@ -60,10 +65,14 @@ from backend.copilot.tools.models import (
|
||||
UnderstandingUpdatedResponse,
|
||||
)
|
||||
from backend.copilot.tracking import track_user_message
|
||||
from backend.data.credit import UsageTransactionMetadata, get_user_credit_model
|
||||
from backend.data.redis_client import get_redis_async
|
||||
from backend.data.understanding import get_business_understanding
|
||||
from backend.data.workspace import get_or_create_workspace
|
||||
from backend.util.exceptions import NotFoundError
|
||||
from backend.util.exceptions import InsufficientBalanceError, NotFoundError
|
||||
from backend.util.settings import Settings
|
||||
|
||||
settings = Settings()
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -432,6 +441,178 @@ async def get_copilot_usage(
|
||||
user_id=user_id,
|
||||
daily_token_limit=daily_limit,
|
||||
weekly_token_limit=weekly_limit,
|
||||
rate_limit_reset_cost=config.rate_limit_reset_cost,
|
||||
)
|
||||
|
||||
|
||||
class RateLimitResetResponse(BaseModel):
|
||||
"""Response from resetting the daily rate limit."""
|
||||
|
||||
success: bool
|
||||
credits_charged: int = Field(description="Credits charged (in cents)")
|
||||
remaining_balance: int = Field(description="Credit balance after charge (in cents)")
|
||||
usage: CoPilotUsageStatus = Field(description="Updated usage status after reset")
|
||||
|
||||
|
||||
@router.post(
|
||||
"/usage/reset",
|
||||
status_code=200,
|
||||
responses={
|
||||
400: {
|
||||
"description": "Bad Request (feature disabled or daily limit not reached)"
|
||||
},
|
||||
402: {"description": "Payment Required (insufficient credits)"},
|
||||
429: {
|
||||
"description": "Too Many Requests (max daily resets exceeded or reset in progress)"
|
||||
},
|
||||
503: {
|
||||
"description": "Service Unavailable (Redis reset failed; credits refunded or support needed)"
|
||||
},
|
||||
},
|
||||
)
|
||||
async def reset_copilot_usage(
|
||||
user_id: Annotated[str, Security(auth.get_user_id)],
|
||||
) -> RateLimitResetResponse:
|
||||
"""Reset the daily CoPilot rate limit by spending credits.
|
||||
|
||||
Allows users who have hit their daily token limit to spend credits
|
||||
to reset their daily usage counter and continue working.
|
||||
Returns 400 if the feature is disabled or the user is not over the limit.
|
||||
Returns 402 if the user has insufficient credits.
|
||||
"""
|
||||
cost = config.rate_limit_reset_cost
|
||||
if cost <= 0:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Rate limit reset is not available.",
|
||||
)
|
||||
|
||||
if not settings.config.enable_credit:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Rate limit reset is not available (credit system is disabled).",
|
||||
)
|
||||
|
||||
daily_limit, weekly_limit = await get_global_rate_limits(
|
||||
user_id, config.daily_token_limit, config.weekly_token_limit
|
||||
)
|
||||
|
||||
if daily_limit <= 0:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="No daily limit is configured — nothing to reset.",
|
||||
)
|
||||
|
||||
# Check max daily resets. get_daily_reset_count returns None when Redis
|
||||
# is unavailable; reject the reset in that case to prevent unlimited
|
||||
# free resets when the counter store is down.
|
||||
reset_count = await get_daily_reset_count(user_id)
|
||||
if reset_count is None:
|
||||
raise HTTPException(
|
||||
status_code=503,
|
||||
detail="Unable to verify reset eligibility — please try again later.",
|
||||
)
|
||||
if config.max_daily_resets > 0 and reset_count >= config.max_daily_resets:
|
||||
raise HTTPException(
|
||||
status_code=429,
|
||||
detail=f"You've used all {config.max_daily_resets} resets for today.",
|
||||
)
|
||||
|
||||
# Acquire a per-user lock to prevent TOCTOU races (concurrent resets).
|
||||
if not await acquire_reset_lock(user_id):
|
||||
raise HTTPException(
|
||||
status_code=429,
|
||||
detail="A reset is already in progress. Please try again.",
|
||||
)
|
||||
|
||||
try:
|
||||
# Verify the user is actually at or over their daily limit.
|
||||
usage_status = await get_usage_status(
|
||||
user_id=user_id,
|
||||
daily_token_limit=daily_limit,
|
||||
weekly_token_limit=weekly_limit,
|
||||
)
|
||||
if daily_limit > 0 and usage_status.daily.used < daily_limit:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="You have not reached your daily limit yet.",
|
||||
)
|
||||
|
||||
# If the weekly limit is also exhausted, resetting the daily counter
|
||||
# won't help — the user would still be blocked by the weekly limit.
|
||||
if weekly_limit > 0 and usage_status.weekly.used >= weekly_limit:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Your weekly limit is also reached. Resetting the daily limit won't help.",
|
||||
)
|
||||
|
||||
# Charge credits.
|
||||
credit_model = await get_user_credit_model(user_id)
|
||||
try:
|
||||
remaining = await credit_model.spend_credits(
|
||||
user_id=user_id,
|
||||
cost=cost,
|
||||
metadata=UsageTransactionMetadata(
|
||||
reason="CoPilot daily rate limit reset",
|
||||
),
|
||||
)
|
||||
except InsufficientBalanceError as e:
|
||||
raise HTTPException(
|
||||
status_code=402,
|
||||
detail="Insufficient credits to reset your rate limit.",
|
||||
) from e
|
||||
|
||||
# Reset daily usage in Redis. If this fails, refund the credits
|
||||
# so the user is not charged for a service they did not receive.
|
||||
if not await reset_daily_usage(user_id, daily_token_limit=daily_limit):
|
||||
# Compensate: refund the charged credits.
|
||||
refunded = False
|
||||
try:
|
||||
await credit_model.top_up_credits(user_id, cost)
|
||||
refunded = True
|
||||
logger.warning(
|
||||
"Refunded %d credits to user %s after Redis reset failure",
|
||||
cost,
|
||||
user_id[:8],
|
||||
)
|
||||
except Exception:
|
||||
logger.error(
|
||||
"CRITICAL: Failed to refund %d credits to user %s "
|
||||
"after Redis reset failure — manual intervention required",
|
||||
cost,
|
||||
user_id[:8],
|
||||
exc_info=True,
|
||||
)
|
||||
if refunded:
|
||||
raise HTTPException(
|
||||
status_code=503,
|
||||
detail="Rate limit reset failed — please try again later. "
|
||||
"Your credits have not been charged.",
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code=503,
|
||||
detail="Rate limit reset failed and the automatic refund "
|
||||
"also failed. Please contact support for assistance.",
|
||||
)
|
||||
|
||||
# Track the reset count for daily cap enforcement.
|
||||
await increment_daily_reset_count(user_id)
|
||||
finally:
|
||||
await release_reset_lock(user_id)
|
||||
|
||||
# Return updated usage status.
|
||||
updated_usage = await get_usage_status(
|
||||
user_id=user_id,
|
||||
daily_token_limit=daily_limit,
|
||||
weekly_token_limit=weekly_limit,
|
||||
rate_limit_reset_cost=config.rate_limit_reset_cost,
|
||||
)
|
||||
|
||||
return RateLimitResetResponse(
|
||||
success=True,
|
||||
credits_charged=cost,
|
||||
remaining_balance=remaining,
|
||||
usage=updated_usage,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -368,6 +368,7 @@ def test_usage_returns_daily_and_weekly(
|
||||
user_id=test_user_id,
|
||||
daily_token_limit=10000,
|
||||
weekly_token_limit=50000,
|
||||
rate_limit_reset_cost=chat_routes.config.rate_limit_reset_cost,
|
||||
)
|
||||
|
||||
|
||||
@@ -380,6 +381,7 @@ def test_usage_uses_config_limits(
|
||||
|
||||
mocker.patch.object(chat_routes.config, "daily_token_limit", 99999)
|
||||
mocker.patch.object(chat_routes.config, "weekly_token_limit", 77777)
|
||||
mocker.patch.object(chat_routes.config, "rate_limit_reset_cost", 500)
|
||||
|
||||
response = client.get("/usage")
|
||||
|
||||
@@ -388,6 +390,7 @@ def test_usage_uses_config_limits(
|
||||
user_id=test_user_id,
|
||||
daily_token_limit=99999,
|
||||
weekly_token_limit=77777,
|
||||
rate_limit_reset_cost=500,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -91,6 +91,20 @@ class ChatConfig(BaseSettings):
|
||||
description="Max tokens per week, resets Monday 00:00 UTC (0 = unlimited)",
|
||||
)
|
||||
|
||||
# Cost (in credits / cents) to reset the daily rate limit using credits.
|
||||
# When a user hits their daily limit, they can spend this amount to reset
|
||||
# the daily counter and keep working. Set to 0 to disable the feature.
|
||||
rate_limit_reset_cost: int = Field(
|
||||
default=500,
|
||||
ge=0,
|
||||
description="Credit cost (in cents) for resetting the daily rate limit. 0 = disabled.",
|
||||
)
|
||||
max_daily_resets: int = Field(
|
||||
default=5,
|
||||
ge=0,
|
||||
description="Maximum number of credit-based rate limit resets per user per day. 0 = unlimited.",
|
||||
)
|
||||
|
||||
# Claude Agent SDK Configuration
|
||||
use_claude_agent_sdk: bool = Field(
|
||||
default=True,
|
||||
|
||||
@@ -36,6 +36,10 @@ class CoPilotUsageStatus(BaseModel):
|
||||
|
||||
daily: UsageWindow
|
||||
weekly: UsageWindow
|
||||
reset_cost: int = Field(
|
||||
default=0,
|
||||
description="Credit cost (in cents) to reset the daily limit. 0 = feature disabled.",
|
||||
)
|
||||
|
||||
|
||||
class RateLimitExceeded(Exception):
|
||||
@@ -61,6 +65,7 @@ async def get_usage_status(
|
||||
user_id: str,
|
||||
daily_token_limit: int,
|
||||
weekly_token_limit: int,
|
||||
rate_limit_reset_cost: int = 0,
|
||||
) -> CoPilotUsageStatus:
|
||||
"""Get current usage status for a user.
|
||||
|
||||
@@ -68,6 +73,7 @@ async def get_usage_status(
|
||||
user_id: The user's ID.
|
||||
daily_token_limit: Max tokens per day (0 = unlimited).
|
||||
weekly_token_limit: Max tokens per week (0 = unlimited).
|
||||
rate_limit_reset_cost: Credit cost (cents) to reset daily limit (0 = disabled).
|
||||
|
||||
Returns:
|
||||
CoPilotUsageStatus with current usage and limits.
|
||||
@@ -97,6 +103,7 @@ async def get_usage_status(
|
||||
limit=weekly_token_limit,
|
||||
resets_at=_weekly_reset_time(now=now),
|
||||
),
|
||||
reset_cost=rate_limit_reset_cost,
|
||||
)
|
||||
|
||||
|
||||
@@ -141,6 +148,110 @@ async def check_rate_limit(
|
||||
raise RateLimitExceeded("weekly", _weekly_reset_time(now=now))
|
||||
|
||||
|
||||
async def reset_daily_usage(user_id: str, daily_token_limit: int = 0) -> bool:
|
||||
"""Reset a user's daily token usage counter in Redis.
|
||||
|
||||
Called after a user pays credits to extend their daily limit.
|
||||
Also reduces the weekly usage counter by ``daily_token_limit`` tokens
|
||||
(clamped to 0) so the user effectively gets one extra day's worth of
|
||||
weekly capacity.
|
||||
|
||||
Args:
|
||||
user_id: The user's ID.
|
||||
daily_token_limit: The configured daily token limit. When positive,
|
||||
the weekly counter is reduced by this amount.
|
||||
|
||||
Fails open: returns False if Redis is unavailable (consistent with
|
||||
the fail-open design of this module).
|
||||
"""
|
||||
now = datetime.now(UTC)
|
||||
try:
|
||||
redis = await get_redis_async()
|
||||
|
||||
# Use a MULTI/EXEC transaction so that DELETE (daily) and DECRBY
|
||||
# (weekly) either both execute or neither does. This prevents the
|
||||
# scenario where the daily counter is cleared but the weekly
|
||||
# counter is not decremented — which would let the caller refund
|
||||
# credits even though the daily limit was already reset.
|
||||
d_key = _daily_key(user_id, now=now)
|
||||
w_key = _weekly_key(user_id, now=now) if daily_token_limit > 0 else None
|
||||
|
||||
pipe = redis.pipeline(transaction=True)
|
||||
pipe.delete(d_key)
|
||||
if w_key is not None:
|
||||
pipe.decrby(w_key, daily_token_limit)
|
||||
results = await pipe.execute()
|
||||
|
||||
# Clamp negative weekly counter to 0 (best-effort; not critical).
|
||||
if w_key is not None:
|
||||
new_val = results[1] # DECRBY result
|
||||
if new_val < 0:
|
||||
await redis.set(w_key, 0, keepttl=True)
|
||||
|
||||
logger.info("Reset daily usage for user %s", user_id[:8])
|
||||
return True
|
||||
except (RedisError, ConnectionError, OSError):
|
||||
logger.warning("Redis unavailable for resetting daily usage")
|
||||
return False
|
||||
|
||||
|
||||
_RESET_LOCK_PREFIX = "copilot:reset_lock"
|
||||
_RESET_COUNT_PREFIX = "copilot:reset_count"
|
||||
|
||||
|
||||
async def acquire_reset_lock(user_id: str, ttl_seconds: int = 10) -> bool:
|
||||
"""Acquire a short-lived lock to serialize rate limit resets per user."""
|
||||
try:
|
||||
redis = await get_redis_async()
|
||||
key = f"{_RESET_LOCK_PREFIX}:{user_id}"
|
||||
return bool(await redis.set(key, "1", nx=True, ex=ttl_seconds))
|
||||
except (RedisError, ConnectionError, OSError) as exc:
|
||||
logger.warning("Redis unavailable for reset lock, rejecting reset: %s", exc)
|
||||
return False
|
||||
|
||||
|
||||
async def release_reset_lock(user_id: str) -> None:
|
||||
"""Release the per-user reset lock."""
|
||||
try:
|
||||
redis = await get_redis_async()
|
||||
await redis.delete(f"{_RESET_LOCK_PREFIX}:{user_id}")
|
||||
except (RedisError, ConnectionError, OSError):
|
||||
pass # Lock will expire via TTL
|
||||
|
||||
|
||||
async def get_daily_reset_count(user_id: str) -> int | None:
|
||||
"""Get how many times the user has reset today.
|
||||
|
||||
Returns None when Redis is unavailable so callers can fail-closed
|
||||
for billed operations (as opposed to failing open for read-only
|
||||
rate-limit checks).
|
||||
"""
|
||||
now = datetime.now(UTC)
|
||||
try:
|
||||
redis = await get_redis_async()
|
||||
key = f"{_RESET_COUNT_PREFIX}:{user_id}:{now.strftime('%Y-%m-%d')}"
|
||||
val = await redis.get(key)
|
||||
return int(val or 0)
|
||||
except (RedisError, ConnectionError, OSError):
|
||||
logger.warning("Redis unavailable for reading daily reset count")
|
||||
return None
|
||||
|
||||
|
||||
async def increment_daily_reset_count(user_id: str) -> None:
|
||||
"""Increment and track how many resets this user has done today."""
|
||||
now = datetime.now(UTC)
|
||||
try:
|
||||
redis = await get_redis_async()
|
||||
key = f"{_RESET_COUNT_PREFIX}:{user_id}:{now.strftime('%Y-%m-%d')}"
|
||||
pipe = redis.pipeline(transaction=True)
|
||||
pipe.incr(key)
|
||||
seconds_until_reset = int((_daily_reset_time(now=now) - now).total_seconds())
|
||||
pipe.expire(key, max(seconds_until_reset, 1))
|
||||
await pipe.execute()
|
||||
except (RedisError, ConnectionError, OSError):
|
||||
logger.warning("Redis unavailable for tracking reset count")
|
||||
|
||||
|
||||
async def record_token_usage(
|
||||
user_id: str,
|
||||
prompt_tokens: int,
|
||||
|
||||
@@ -12,6 +12,7 @@ from .rate_limit import (
|
||||
check_rate_limit,
|
||||
get_usage_status,
|
||||
record_token_usage,
|
||||
reset_daily_usage,
|
||||
)
|
||||
|
||||
_USER = "test-user-rl"
|
||||
@@ -332,3 +333,91 @@ class TestRecordTokenUsage:
|
||||
):
|
||||
# Should not raise — fail-open
|
||||
await record_token_usage(_USER, prompt_tokens=100, completion_tokens=50)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# reset_daily_usage
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestResetDailyUsage:
|
||||
@staticmethod
|
||||
def _make_pipeline_mock(decrby_result: int = 0) -> MagicMock:
|
||||
"""Create a pipeline mock that returns [delete_result, decrby_result]."""
|
||||
pipe = MagicMock()
|
||||
pipe.execute = AsyncMock(return_value=[1, decrby_result])
|
||||
return pipe
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deletes_daily_key(self):
|
||||
mock_pipe = self._make_pipeline_mock(decrby_result=0)
|
||||
mock_redis = AsyncMock()
|
||||
mock_redis.pipeline = lambda **_kw: mock_pipe
|
||||
|
||||
with patch(
|
||||
"backend.copilot.rate_limit.get_redis_async",
|
||||
return_value=mock_redis,
|
||||
):
|
||||
result = await reset_daily_usage(_USER, daily_token_limit=10000)
|
||||
|
||||
assert result is True
|
||||
mock_pipe.delete.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reduces_weekly_usage_via_decrby(self):
|
||||
"""Weekly counter should be reduced via DECRBY in the pipeline."""
|
||||
mock_pipe = self._make_pipeline_mock(decrby_result=35000)
|
||||
mock_redis = AsyncMock()
|
||||
mock_redis.pipeline = lambda **_kw: mock_pipe
|
||||
|
||||
with patch(
|
||||
"backend.copilot.rate_limit.get_redis_async",
|
||||
return_value=mock_redis,
|
||||
):
|
||||
await reset_daily_usage(_USER, daily_token_limit=10000)
|
||||
|
||||
mock_pipe.decrby.assert_called_once()
|
||||
mock_redis.set.assert_not_called() # 35000 > 0, no clamp needed
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_clamps_negative_weekly_to_zero(self):
|
||||
"""If DECRBY goes negative, SET to 0 (outside the pipeline)."""
|
||||
mock_pipe = self._make_pipeline_mock(decrby_result=-5000)
|
||||
mock_redis = AsyncMock()
|
||||
mock_redis.pipeline = lambda **_kw: mock_pipe
|
||||
|
||||
with patch(
|
||||
"backend.copilot.rate_limit.get_redis_async",
|
||||
return_value=mock_redis,
|
||||
):
|
||||
await reset_daily_usage(_USER, daily_token_limit=10000)
|
||||
|
||||
mock_pipe.decrby.assert_called_once()
|
||||
mock_redis.set.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_weekly_reduction_when_daily_limit_zero(self):
|
||||
"""When daily_token_limit is 0, weekly counter should not be touched."""
|
||||
mock_pipe = self._make_pipeline_mock()
|
||||
mock_pipe.execute = AsyncMock(return_value=[1]) # only delete result
|
||||
mock_redis = AsyncMock()
|
||||
mock_redis.pipeline = lambda **_kw: mock_pipe
|
||||
|
||||
with patch(
|
||||
"backend.copilot.rate_limit.get_redis_async",
|
||||
return_value=mock_redis,
|
||||
):
|
||||
await reset_daily_usage(_USER, daily_token_limit=0)
|
||||
|
||||
mock_pipe.delete.assert_called_once()
|
||||
mock_pipe.decrby.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_false_when_redis_unavailable(self):
|
||||
with patch(
|
||||
"backend.copilot.rate_limit.get_redis_async",
|
||||
side_effect=ConnectionError("Redis down"),
|
||||
):
|
||||
result = await reset_daily_usage(_USER, daily_token_limit=10000)
|
||||
|
||||
assert result is False
|
||||
|
||||
294
autogpt_platform/backend/backend/copilot/reset_usage_test.py
Normal file
294
autogpt_platform/backend/backend/copilot/reset_usage_test.py
Normal file
@@ -0,0 +1,294 @@
|
||||
"""Unit tests for the POST /usage/reset endpoint."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
||||
from backend.api.features.chat.routes import reset_copilot_usage
|
||||
from backend.copilot.rate_limit import CoPilotUsageStatus, UsageWindow
|
||||
from backend.util.exceptions import InsufficientBalanceError
|
||||
|
||||
|
||||
# Minimal config mock matching ChatConfig fields used by the endpoint.
|
||||
def _make_config(
|
||||
rate_limit_reset_cost: int = 500,
|
||||
daily_token_limit: int = 2_500_000,
|
||||
weekly_token_limit: int = 12_500_000,
|
||||
max_daily_resets: int = 5,
|
||||
):
|
||||
cfg = MagicMock()
|
||||
cfg.rate_limit_reset_cost = rate_limit_reset_cost
|
||||
cfg.daily_token_limit = daily_token_limit
|
||||
cfg.weekly_token_limit = weekly_token_limit
|
||||
cfg.max_daily_resets = max_daily_resets
|
||||
return cfg
|
||||
|
||||
|
||||
def _usage(daily_used: int = 3_000_000, daily_limit: int = 2_500_000):
|
||||
return CoPilotUsageStatus(
|
||||
daily=UsageWindow(
|
||||
used=daily_used,
|
||||
limit=daily_limit,
|
||||
resets_at=datetime.now(UTC) + timedelta(hours=6),
|
||||
),
|
||||
weekly=UsageWindow(
|
||||
used=5_000_000,
|
||||
limit=12_500_000,
|
||||
resets_at=datetime.now(UTC) + timedelta(days=3),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
_MODULE = "backend.api.features.chat.routes"
|
||||
|
||||
|
||||
def _mock_settings(enable_credit: bool = True):
|
||||
"""Return a mock Settings object with the given enable_credit flag."""
|
||||
mock = MagicMock()
|
||||
mock.config.enable_credit = enable_credit
|
||||
return mock
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestResetCopilotUsage:
|
||||
async def test_feature_disabled_returns_400(self):
|
||||
"""When rate_limit_reset_cost=0, endpoint returns 400."""
|
||||
|
||||
with patch(f"{_MODULE}.config", _make_config(rate_limit_reset_cost=0)):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await reset_copilot_usage(user_id="user-1")
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "not available" in exc_info.value.detail
|
||||
|
||||
async def test_no_daily_limit_returns_400(self):
|
||||
"""When daily_token_limit=0 (unlimited), endpoint returns 400."""
|
||||
|
||||
with (
|
||||
patch(f"{_MODULE}.config", _make_config(daily_token_limit=0)),
|
||||
patch(f"{_MODULE}.settings", _mock_settings()),
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await reset_copilot_usage(user_id="user-1")
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "nothing to reset" in exc_info.value.detail.lower()
|
||||
|
||||
async def test_not_at_limit_returns_400(self):
|
||||
"""When user hasn't hit their daily limit, returns 400."""
|
||||
|
||||
cfg = _make_config()
|
||||
with (
|
||||
patch(f"{_MODULE}.config", cfg),
|
||||
patch(f"{_MODULE}.settings", _mock_settings()),
|
||||
patch(f"{_MODULE}.get_daily_reset_count", AsyncMock(return_value=0)),
|
||||
patch(f"{_MODULE}.acquire_reset_lock", AsyncMock(return_value=True)),
|
||||
patch(f"{_MODULE}.release_reset_lock", AsyncMock()) as mock_release,
|
||||
patch(
|
||||
f"{_MODULE}.get_usage_status",
|
||||
AsyncMock(return_value=_usage(daily_used=1_000_000)),
|
||||
),
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await reset_copilot_usage(user_id="user-1")
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "not reached" in exc_info.value.detail
|
||||
mock_release.assert_awaited_once()
|
||||
|
||||
async def test_insufficient_credits_returns_402(self):
|
||||
"""When user doesn't have enough credits, returns 402."""
|
||||
|
||||
mock_credit_model = AsyncMock()
|
||||
mock_credit_model.spend_credits.side_effect = InsufficientBalanceError(
|
||||
message="Insufficient balance",
|
||||
user_id="user-1",
|
||||
balance=50,
|
||||
amount=200,
|
||||
)
|
||||
|
||||
cfg = _make_config()
|
||||
with (
|
||||
patch(f"{_MODULE}.config", cfg),
|
||||
patch(f"{_MODULE}.settings", _mock_settings()),
|
||||
patch(f"{_MODULE}.get_daily_reset_count", AsyncMock(return_value=0)),
|
||||
patch(f"{_MODULE}.acquire_reset_lock", AsyncMock(return_value=True)),
|
||||
patch(f"{_MODULE}.release_reset_lock", AsyncMock()) as mock_release,
|
||||
patch(
|
||||
f"{_MODULE}.get_usage_status",
|
||||
AsyncMock(return_value=_usage()),
|
||||
),
|
||||
patch(
|
||||
f"{_MODULE}.get_user_credit_model",
|
||||
AsyncMock(return_value=mock_credit_model),
|
||||
),
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await reset_copilot_usage(user_id="user-1")
|
||||
assert exc_info.value.status_code == 402
|
||||
mock_release.assert_awaited_once()
|
||||
|
||||
async def test_happy_path(self):
|
||||
"""Successful reset: charges credits, resets usage, returns response."""
|
||||
|
||||
mock_credit_model = AsyncMock()
|
||||
mock_credit_model.spend_credits.return_value = 1500 # remaining balance
|
||||
|
||||
cfg = _make_config()
|
||||
updated_usage = _usage(daily_used=0)
|
||||
|
||||
with (
|
||||
patch(f"{_MODULE}.config", cfg),
|
||||
patch(f"{_MODULE}.settings", _mock_settings()),
|
||||
patch(f"{_MODULE}.get_daily_reset_count", AsyncMock(return_value=0)),
|
||||
patch(f"{_MODULE}.acquire_reset_lock", AsyncMock(return_value=True)),
|
||||
patch(f"{_MODULE}.release_reset_lock", AsyncMock()),
|
||||
patch(
|
||||
f"{_MODULE}.get_usage_status",
|
||||
AsyncMock(side_effect=[_usage(), updated_usage]),
|
||||
),
|
||||
patch(
|
||||
f"{_MODULE}.get_user_credit_model",
|
||||
AsyncMock(return_value=mock_credit_model),
|
||||
),
|
||||
patch(
|
||||
f"{_MODULE}.reset_daily_usage", AsyncMock(return_value=True)
|
||||
) as mock_reset,
|
||||
patch(f"{_MODULE}.increment_daily_reset_count", AsyncMock()) as mock_incr,
|
||||
):
|
||||
result = await reset_copilot_usage(user_id="user-1")
|
||||
assert result.success is True
|
||||
assert result.credits_charged == 500
|
||||
assert result.remaining_balance == 1500
|
||||
mock_reset.assert_awaited_once()
|
||||
mock_incr.assert_awaited_once()
|
||||
|
||||
async def test_max_daily_resets_exceeded(self):
|
||||
"""When user has exhausted daily resets, returns 429."""
|
||||
|
||||
cfg = _make_config(max_daily_resets=3)
|
||||
with (
|
||||
patch(f"{_MODULE}.config", cfg),
|
||||
patch(f"{_MODULE}.settings", _mock_settings()),
|
||||
patch(f"{_MODULE}.get_daily_reset_count", AsyncMock(return_value=3)),
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await reset_copilot_usage(user_id="user-1")
|
||||
assert exc_info.value.status_code == 429
|
||||
|
||||
async def test_credit_system_disabled_returns_400(self):
|
||||
"""When enable_credit=False, endpoint returns 400."""
|
||||
|
||||
with (
|
||||
patch(f"{_MODULE}.config", _make_config()),
|
||||
patch(f"{_MODULE}.settings", _mock_settings(enable_credit=False)),
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await reset_copilot_usage(user_id="user-1")
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "credit system is disabled" in exc_info.value.detail.lower()
|
||||
|
||||
async def test_weekly_limit_exhausted_returns_400(self):
|
||||
"""When the weekly limit is also exhausted, resetting daily won't help."""
|
||||
|
||||
cfg = _make_config()
|
||||
weekly_exhausted = CoPilotUsageStatus(
|
||||
daily=UsageWindow(
|
||||
used=3_000_000,
|
||||
limit=2_500_000,
|
||||
resets_at=datetime.now(UTC) + timedelta(hours=6),
|
||||
),
|
||||
weekly=UsageWindow(
|
||||
used=12_500_000,
|
||||
limit=12_500_000,
|
||||
resets_at=datetime.now(UTC) + timedelta(days=3),
|
||||
),
|
||||
)
|
||||
with (
|
||||
patch(f"{_MODULE}.config", cfg),
|
||||
patch(f"{_MODULE}.settings", _mock_settings()),
|
||||
patch(f"{_MODULE}.get_daily_reset_count", AsyncMock(return_value=0)),
|
||||
patch(f"{_MODULE}.acquire_reset_lock", AsyncMock(return_value=True)),
|
||||
patch(f"{_MODULE}.release_reset_lock", AsyncMock()) as mock_release,
|
||||
patch(
|
||||
f"{_MODULE}.get_usage_status",
|
||||
AsyncMock(return_value=weekly_exhausted),
|
||||
),
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await reset_copilot_usage(user_id="user-1")
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "weekly" in exc_info.value.detail.lower()
|
||||
mock_release.assert_awaited_once()
|
||||
|
||||
async def test_redis_failure_for_reset_count_returns_503(self):
|
||||
"""When Redis is unavailable for get_daily_reset_count, returns 503."""
|
||||
|
||||
with (
|
||||
patch(f"{_MODULE}.config", _make_config()),
|
||||
patch(f"{_MODULE}.settings", _mock_settings()),
|
||||
patch(f"{_MODULE}.get_daily_reset_count", AsyncMock(return_value=None)),
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await reset_copilot_usage(user_id="user-1")
|
||||
assert exc_info.value.status_code == 503
|
||||
assert "verify" in exc_info.value.detail.lower()
|
||||
|
||||
async def test_redis_reset_failure_refunds_credits(self):
|
||||
"""When reset_daily_usage fails, credits are refunded and 503 returned."""
|
||||
|
||||
mock_credit_model = AsyncMock()
|
||||
mock_credit_model.spend_credits.return_value = 1500
|
||||
|
||||
cfg = _make_config()
|
||||
with (
|
||||
patch(f"{_MODULE}.config", cfg),
|
||||
patch(f"{_MODULE}.settings", _mock_settings()),
|
||||
patch(f"{_MODULE}.get_daily_reset_count", AsyncMock(return_value=0)),
|
||||
patch(f"{_MODULE}.acquire_reset_lock", AsyncMock(return_value=True)),
|
||||
patch(f"{_MODULE}.release_reset_lock", AsyncMock()),
|
||||
patch(
|
||||
f"{_MODULE}.get_usage_status",
|
||||
AsyncMock(return_value=_usage()),
|
||||
),
|
||||
patch(
|
||||
f"{_MODULE}.get_user_credit_model",
|
||||
AsyncMock(return_value=mock_credit_model),
|
||||
),
|
||||
patch(f"{_MODULE}.reset_daily_usage", AsyncMock(return_value=False)),
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await reset_copilot_usage(user_id="user-1")
|
||||
assert exc_info.value.status_code == 503
|
||||
assert "not been charged" in exc_info.value.detail
|
||||
mock_credit_model.top_up_credits.assert_awaited_once()
|
||||
|
||||
async def test_redis_reset_failure_refund_also_fails(self):
|
||||
"""When both reset and refund fail, error message reflects the truth."""
|
||||
|
||||
mock_credit_model = AsyncMock()
|
||||
mock_credit_model.spend_credits.return_value = 1500
|
||||
mock_credit_model.top_up_credits.side_effect = RuntimeError("db down")
|
||||
|
||||
cfg = _make_config()
|
||||
with (
|
||||
patch(f"{_MODULE}.config", cfg),
|
||||
patch(f"{_MODULE}.settings", _mock_settings()),
|
||||
patch(f"{_MODULE}.get_daily_reset_count", AsyncMock(return_value=0)),
|
||||
patch(f"{_MODULE}.acquire_reset_lock", AsyncMock(return_value=True)),
|
||||
patch(f"{_MODULE}.release_reset_lock", AsyncMock()),
|
||||
patch(
|
||||
f"{_MODULE}.get_usage_status",
|
||||
AsyncMock(return_value=_usage()),
|
||||
),
|
||||
patch(
|
||||
f"{_MODULE}.get_user_credit_model",
|
||||
AsyncMock(return_value=mock_credit_model),
|
||||
),
|
||||
patch(f"{_MODULE}.reset_daily_usage", AsyncMock(return_value=False)),
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await reset_copilot_usage(user_id="user-1")
|
||||
assert exc_info.value.status_code == 503
|
||||
assert "contact support" in exc_info.value.detail.lower()
|
||||
@@ -1,9 +1,14 @@
|
||||
"use client";
|
||||
|
||||
import type { CoPilotUsageStatus } from "@/app/api/__generated__/models/coPilotUsageStatus";
|
||||
import { useGetV2GetCopilotUsage } from "@/app/api/__generated__/endpoints/chat/chat";
|
||||
import { toast } from "@/components/molecules/Toast/use-toast";
|
||||
import useCredits from "@/hooks/useCredits";
|
||||
import { Flag, useGetFlag } from "@/services/feature-flags/use-get-flag";
|
||||
import { SidebarProvider } from "@/components/ui/sidebar";
|
||||
import { cn } from "@/lib/utils";
|
||||
import { UploadSimple } from "@phosphor-icons/react";
|
||||
import { useCallback, useRef, useState } from "react";
|
||||
import { useCallback, useEffect, useRef, useState } from "react";
|
||||
import { ChatContainer } from "./components/ChatContainer/ChatContainer";
|
||||
import { ChatSidebar } from "./components/ChatSidebar/ChatSidebar";
|
||||
import { DeleteChatDialog } from "./components/DeleteChatDialog/DeleteChatDialog";
|
||||
@@ -11,6 +16,7 @@ import { MobileDrawer } from "./components/MobileDrawer/MobileDrawer";
|
||||
import { MobileHeader } from "./components/MobileHeader/MobileHeader";
|
||||
import { NotificationBanner } from "./components/NotificationBanner/NotificationBanner";
|
||||
import { NotificationDialog } from "./components/NotificationDialog/NotificationDialog";
|
||||
import { RateLimitResetDialog } from "./components/RateLimitResetDialog/RateLimitResetDialog";
|
||||
import { ScaleLoader } from "./components/ScaleLoader/ScaleLoader";
|
||||
import { useCopilotPage } from "./useCopilotPage";
|
||||
|
||||
@@ -89,8 +95,45 @@ export function CopilotPage() {
|
||||
isDeleting,
|
||||
handleConfirmDelete,
|
||||
handleCancelDelete,
|
||||
// Rate limit reset
|
||||
rateLimitMessage,
|
||||
dismissRateLimit,
|
||||
} = useCopilotPage();
|
||||
|
||||
const {
|
||||
data: usage,
|
||||
isSuccess: hasUsage,
|
||||
isError: usageError,
|
||||
} = useGetV2GetCopilotUsage({
|
||||
query: {
|
||||
select: (res) => res.data as CoPilotUsageStatus,
|
||||
refetchInterval: 30000,
|
||||
staleTime: 10000,
|
||||
},
|
||||
});
|
||||
const resetCost = usage?.reset_cost;
|
||||
|
||||
const isBillingEnabled = useGetFlag(Flag.ENABLE_PLATFORM_PAYMENT);
|
||||
const { credits, fetchCredits } = useCredits({ fetchInitialCredits: true });
|
||||
const hasInsufficientCredits =
|
||||
credits !== null && resetCost != null && credits < resetCost;
|
||||
|
||||
// Fall back to a toast when the credit-based reset feature is disabled or
|
||||
// when the usage query fails (so the user still gets feedback).
|
||||
useEffect(() => {
|
||||
if (
|
||||
rateLimitMessage &&
|
||||
(usageError || (hasUsage && (resetCost ?? 0) <= 0))
|
||||
) {
|
||||
toast({
|
||||
title: "Usage limit reached",
|
||||
description: rateLimitMessage,
|
||||
variant: "destructive",
|
||||
});
|
||||
dismissRateLimit();
|
||||
}
|
||||
}, [rateLimitMessage, resetCost, hasUsage, usageError, dismissRateLimit]);
|
||||
|
||||
if (isUserLoading || !isLoggedIn) {
|
||||
return (
|
||||
<div className="fixed inset-0 z-50 flex items-center justify-center bg-[#f8f8f9]">
|
||||
@@ -168,6 +211,20 @@ export function CopilotPage() {
|
||||
/>
|
||||
)}
|
||||
<NotificationDialog />
|
||||
<RateLimitResetDialog
|
||||
isOpen={!!rateLimitMessage && hasUsage && (resetCost ?? 0) > 0}
|
||||
onClose={dismissRateLimit}
|
||||
resetCost={resetCost ?? 0}
|
||||
resetMessage={rateLimitMessage ?? ""}
|
||||
isWeeklyExhausted={
|
||||
hasUsage &&
|
||||
usage.weekly.limit > 0 &&
|
||||
usage.weekly.used >= usage.weekly.limit
|
||||
}
|
||||
hasInsufficientCredits={hasInsufficientCredits}
|
||||
isBillingEnabled={isBillingEnabled}
|
||||
onCreditChange={fetchCredits}
|
||||
/>
|
||||
</SidebarProvider>
|
||||
);
|
||||
}
|
||||
|
||||
@@ -0,0 +1,114 @@
|
||||
"use client";
|
||||
|
||||
import { Button } from "@/components/atoms/Button/Button";
|
||||
import { Text } from "@/components/atoms/Text/Text";
|
||||
import { Dialog } from "@/components/molecules/Dialog/Dialog";
|
||||
import { useRouter } from "next/navigation";
|
||||
import { useEffect } from "react";
|
||||
import { useResetRateLimit } from "../../hooks/useResetRateLimit";
|
||||
|
||||
interface Props {
|
||||
isOpen: boolean;
|
||||
onClose: () => void;
|
||||
resetCost: number;
|
||||
resetMessage: string;
|
||||
isWeeklyExhausted?: boolean;
|
||||
hasInsufficientCredits?: boolean;
|
||||
isBillingEnabled?: boolean;
|
||||
onCreditChange?: () => void;
|
||||
}
|
||||
|
||||
function formatCents(cents: number): string {
|
||||
return `$${(cents / 100).toFixed(2)}`;
|
||||
}
|
||||
|
||||
export function RateLimitResetDialog({
|
||||
isOpen,
|
||||
onClose,
|
||||
resetCost,
|
||||
resetMessage,
|
||||
isWeeklyExhausted = false,
|
||||
hasInsufficientCredits = false,
|
||||
isBillingEnabled = false,
|
||||
onCreditChange,
|
||||
}: Props) {
|
||||
const { resetUsage, isPending } = useResetRateLimit({
|
||||
onSuccess: onClose,
|
||||
onCreditChange,
|
||||
});
|
||||
const router = useRouter();
|
||||
|
||||
// Refresh the credit balance each time the dialog opens so we never
|
||||
// block a valid reset due to a stale client-side balance.
|
||||
useEffect(() => {
|
||||
if (isOpen) onCreditChange?.();
|
||||
}, [isOpen]); // eslint-disable-line react-hooks/exhaustive-deps
|
||||
|
||||
// Whether to hide the reset button entirely
|
||||
const cannotReset = isWeeklyExhausted || hasInsufficientCredits;
|
||||
|
||||
return (
|
||||
<Dialog
|
||||
title="Usage limit reached"
|
||||
styling={{ maxWidth: "28rem", minWidth: "auto" }}
|
||||
controlled={{
|
||||
isOpen,
|
||||
set: async (open) => {
|
||||
if (!open) onClose();
|
||||
},
|
||||
}}
|
||||
>
|
||||
<Dialog.Content>
|
||||
<div className="flex flex-col gap-3">
|
||||
<Text variant="body">{resetMessage}</Text>
|
||||
{isWeeklyExhausted ? (
|
||||
<Text variant="body">
|
||||
Your weekly limit is also reached, so resetting the daily limit
|
||||
won't help. Please wait for your limits to reset.
|
||||
</Text>
|
||||
) : hasInsufficientCredits ? (
|
||||
<Text variant="body">
|
||||
You don't have enough credits to reset your daily limit.
|
||||
{isBillingEnabled
|
||||
? " Add credits to continue working."
|
||||
: " Please wait for your limits to reset."}
|
||||
</Text>
|
||||
) : (
|
||||
<Text variant="body">
|
||||
You can spend{" "}
|
||||
<Text variant="body-medium" as="span">
|
||||
{formatCents(resetCost)}
|
||||
</Text>{" "}
|
||||
in credits to reset your daily limit and continue working.
|
||||
</Text>
|
||||
)}
|
||||
</div>
|
||||
<Dialog.Footer className="!justify-center">
|
||||
<Button variant="secondary" onClick={onClose} disabled={isPending}>
|
||||
{cannotReset ? "OK" : "Wait for reset"}
|
||||
</Button>
|
||||
{hasInsufficientCredits && isBillingEnabled && (
|
||||
<Button
|
||||
variant="primary"
|
||||
onClick={() => {
|
||||
onClose();
|
||||
router.push("/profile/credits");
|
||||
}}
|
||||
>
|
||||
Add credits
|
||||
</Button>
|
||||
)}
|
||||
{!cannotReset && (
|
||||
<Button
|
||||
variant="primary"
|
||||
onClick={() => resetUsage()}
|
||||
loading={isPending}
|
||||
>
|
||||
Reset for {formatCents(resetCost)}
|
||||
</Button>
|
||||
)}
|
||||
</Dialog.Footer>
|
||||
</Dialog.Content>
|
||||
</Dialog>
|
||||
);
|
||||
}
|
||||
@@ -1,5 +1,7 @@
|
||||
import type { CoPilotUsageStatus } from "@/app/api/__generated__/models/coPilotUsageStatus";
|
||||
import { useGetV2GetCopilotUsage } from "@/app/api/__generated__/endpoints/chat/chat";
|
||||
import useCredits from "@/hooks/useCredits";
|
||||
import { Flag, useGetFlag } from "@/services/feature-flags/use-get-flag";
|
||||
import {
|
||||
Popover,
|
||||
PopoverContent,
|
||||
@@ -20,6 +22,12 @@ export function UsageLimits() {
|
||||
},
|
||||
});
|
||||
|
||||
const isBillingEnabled = useGetFlag(Flag.ENABLE_PLATFORM_PAYMENT);
|
||||
const { credits, fetchCredits } = useCredits({ fetchInitialCredits: true });
|
||||
const resetCost = usage?.reset_cost;
|
||||
const hasInsufficientCredits =
|
||||
credits !== null && resetCost != null && credits < resetCost;
|
||||
|
||||
if (isLoading || !usage?.daily || !usage?.weekly) return null;
|
||||
if (usage.daily.limit <= 0 && usage.weekly.limit <= 0) return null;
|
||||
|
||||
@@ -31,7 +39,12 @@ export function UsageLimits() {
|
||||
</Button>
|
||||
</PopoverTrigger>
|
||||
<PopoverContent align="start" className="w-64 p-3">
|
||||
<UsagePanelContent usage={usage} />
|
||||
<UsagePanelContent
|
||||
usage={usage}
|
||||
hasInsufficientCredits={hasInsufficientCredits}
|
||||
isBillingEnabled={isBillingEnabled}
|
||||
onCreditChange={fetchCredits}
|
||||
/>
|
||||
</PopoverContent>
|
||||
</Popover>
|
||||
);
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
import type { CoPilotUsageStatus } from "@/app/api/__generated__/models/coPilotUsageStatus";
|
||||
import { Button } from "@/components/atoms/Button/Button";
|
||||
import Link from "next/link";
|
||||
import { useResetRateLimit } from "../../hooks/useResetRateLimit";
|
||||
|
||||
export function formatResetTime(
|
||||
resetsAt: Date | string,
|
||||
@@ -70,15 +72,50 @@ function UsageBar({
|
||||
);
|
||||
}
|
||||
|
||||
function ResetButton({
|
||||
cost,
|
||||
onCreditChange,
|
||||
}: {
|
||||
cost: number;
|
||||
onCreditChange?: () => void;
|
||||
}) {
|
||||
const { resetUsage, isPending } = useResetRateLimit({ onCreditChange });
|
||||
|
||||
return (
|
||||
<Button
|
||||
variant="primary"
|
||||
size="small"
|
||||
onClick={() => resetUsage()}
|
||||
loading={isPending}
|
||||
className="mt-1 w-full text-[11px]"
|
||||
>
|
||||
{isPending
|
||||
? "Resetting..."
|
||||
: `Reset daily limit for $${(cost / 100).toFixed(2)}`}
|
||||
</Button>
|
||||
);
|
||||
}
|
||||
|
||||
export function UsagePanelContent({
|
||||
usage,
|
||||
showBillingLink = true,
|
||||
hasInsufficientCredits = false,
|
||||
isBillingEnabled = false,
|
||||
onCreditChange,
|
||||
}: {
|
||||
usage: CoPilotUsageStatus;
|
||||
showBillingLink?: boolean;
|
||||
hasInsufficientCredits?: boolean;
|
||||
isBillingEnabled?: boolean;
|
||||
onCreditChange?: () => void;
|
||||
}) {
|
||||
const hasDailyLimit = usage.daily.limit > 0;
|
||||
const hasWeeklyLimit = usage.weekly.limit > 0;
|
||||
const isDailyExhausted =
|
||||
hasDailyLimit && usage.daily.used >= usage.daily.limit;
|
||||
const isWeeklyExhausted =
|
||||
hasWeeklyLimit && usage.weekly.used >= usage.weekly.limit;
|
||||
const resetCost = usage.reset_cost ?? 0;
|
||||
|
||||
if (!hasDailyLimit && !hasWeeklyLimit) {
|
||||
return (
|
||||
@@ -105,6 +142,23 @@ export function UsagePanelContent({
|
||||
resetsAt={usage.weekly.resets_at}
|
||||
/>
|
||||
)}
|
||||
{isDailyExhausted &&
|
||||
!isWeeklyExhausted &&
|
||||
resetCost > 0 &&
|
||||
!hasInsufficientCredits && (
|
||||
<ResetButton cost={resetCost} onCreditChange={onCreditChange} />
|
||||
)}
|
||||
{isDailyExhausted &&
|
||||
!isWeeklyExhausted &&
|
||||
hasInsufficientCredits &&
|
||||
isBillingEnabled && (
|
||||
<Link
|
||||
href="/profile/credits"
|
||||
className="mt-1 inline-flex w-full items-center justify-center rounded-md bg-primary px-3 py-1.5 text-[11px] font-medium text-primary-foreground hover:bg-primary/90"
|
||||
>
|
||||
Add credits to reset
|
||||
</Link>
|
||||
)}
|
||||
{showBillingLink && (
|
||||
<Link
|
||||
href="/profile/credits"
|
||||
|
||||
@@ -0,0 +1,48 @@
|
||||
import {
|
||||
usePostV2ResetCopilotUsage,
|
||||
getGetV2GetCopilotUsageQueryKey,
|
||||
} from "@/app/api/__generated__/endpoints/chat/chat";
|
||||
import { toast } from "@/components/molecules/Toast/use-toast";
|
||||
import { ApiError } from "@/lib/autogpt-server-api";
|
||||
import { useQueryClient } from "@tanstack/react-query";
|
||||
|
||||
export function useResetRateLimit(options?: {
|
||||
onSuccess?: () => void;
|
||||
onCreditChange?: () => void;
|
||||
}) {
|
||||
const queryClient = useQueryClient();
|
||||
const { mutate: resetUsage, isPending } = usePostV2ResetCopilotUsage({
|
||||
mutation: {
|
||||
onSuccess: async () => {
|
||||
// Await the usage refetch so the UI shows updated limits before
|
||||
// closing the dialog or re-enabling the reset CTA.
|
||||
// invalidateQueries already triggers a refetch for active queries.
|
||||
await queryClient.invalidateQueries({
|
||||
queryKey: getGetV2GetCopilotUsageQueryKey(),
|
||||
});
|
||||
options?.onCreditChange?.();
|
||||
toast({
|
||||
title: "Rate limit reset",
|
||||
description:
|
||||
"Your daily usage limit has been reset. You can continue working.",
|
||||
});
|
||||
options?.onSuccess?.();
|
||||
},
|
||||
onError: (error: unknown) => {
|
||||
const message =
|
||||
error instanceof ApiError
|
||||
? (error.response?.detail ?? error.message)
|
||||
: error instanceof Error
|
||||
? error.message
|
||||
: "Failed to reset limit.";
|
||||
toast({
|
||||
title: "Reset failed",
|
||||
description: message,
|
||||
variant: "destructive",
|
||||
});
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
return { resetUsage, isPending };
|
||||
}
|
||||
@@ -56,6 +56,8 @@ export function useCopilotPage() {
|
||||
isReconnecting,
|
||||
isSyncing,
|
||||
isUserStoppingRef,
|
||||
rateLimitMessage,
|
||||
dismissRateLimit,
|
||||
} = useCopilotStream({
|
||||
sessionId,
|
||||
hydratedMessages,
|
||||
@@ -375,5 +377,8 @@ export function useCopilotPage() {
|
||||
handleDeleteClick,
|
||||
handleConfirmDelete,
|
||||
handleCancelDelete,
|
||||
// Rate limit reset
|
||||
rateLimitMessage,
|
||||
dismissRateLimit,
|
||||
};
|
||||
}
|
||||
|
||||
@@ -10,7 +10,7 @@ import { useChat } from "@ai-sdk/react";
|
||||
import { useQueryClient } from "@tanstack/react-query";
|
||||
import { DefaultChatTransport } from "ai";
|
||||
import type { FileUIPart, UIMessage } from "ai";
|
||||
import { useEffect, useMemo, useRef, useState } from "react";
|
||||
import { useCallback, useEffect, useMemo, useRef, useState } from "react";
|
||||
import {
|
||||
deduplicateMessages,
|
||||
hasActiveBackendStream,
|
||||
@@ -47,6 +47,8 @@ export function useCopilotStream({
|
||||
refetchSession,
|
||||
}: UseCopilotStreamArgs) {
|
||||
const queryClient = useQueryClient();
|
||||
const [rateLimitMessage, setRateLimitMessage] = useState<string | null>(null);
|
||||
const dismissRateLimit = useCallback(() => setRateLimitMessage(null), []);
|
||||
|
||||
// Connect directly to the Python backend for SSE, bypassing the Next.js
|
||||
// serverless proxy. This eliminates the Vercel 800s function timeout that
|
||||
@@ -196,13 +198,10 @@ export function useCopilotStream({
|
||||
}
|
||||
const isRateLimited = errorDetail.toLowerCase().includes("usage limit");
|
||||
if (isRateLimited) {
|
||||
toast({
|
||||
title: "Usage limit reached",
|
||||
description:
|
||||
errorDetail ||
|
||||
setRateLimitMessage(
|
||||
errorDetail ||
|
||||
"You've reached your usage limit. Please try again later.",
|
||||
variant: "destructive",
|
||||
});
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
@@ -379,6 +378,7 @@ export function useCopilotStream({
|
||||
reconnectAttemptsRef.current = 0;
|
||||
isReconnectScheduledRef.current = false;
|
||||
setIsReconnectScheduled(false);
|
||||
setRateLimitMessage(null);
|
||||
hasShownDisconnectToast.current = false;
|
||||
isUserStoppingRef.current = false;
|
||||
setReconnectExhausted(false);
|
||||
@@ -487,5 +487,7 @@ export function useCopilotStream({
|
||||
isReconnecting,
|
||||
isSyncing,
|
||||
isUserStoppingRef,
|
||||
rateLimitMessage,
|
||||
dismissRateLimit,
|
||||
};
|
||||
}
|
||||
|
||||
@@ -1404,6 +1404,40 @@
|
||||
"security": [{ "HTTPBearerJWT": [] }]
|
||||
}
|
||||
},
|
||||
"/api/chat/usage/reset": {
|
||||
"post": {
|
||||
"tags": ["v2", "chat", "chat"],
|
||||
"summary": "Reset Copilot Usage",
|
||||
"description": "Reset the daily CoPilot rate limit by spending credits.\n\nAllows users who have hit their daily token limit to spend credits\nto reset their daily usage counter and continue working.\nReturns 400 if the feature is disabled or the user is not over the limit.\nReturns 402 if the user has insufficient credits.",
|
||||
"operationId": "postV2ResetCopilotUsage",
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "Successful Response",
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/RateLimitResetResponse"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"400": {
|
||||
"description": "Bad Request (feature disabled or daily limit not reached)"
|
||||
},
|
||||
"401": {
|
||||
"$ref": "#/components/responses/HTTP401NotAuthenticatedError"
|
||||
},
|
||||
"402": { "description": "Payment Required (insufficient credits)" },
|
||||
"429": {
|
||||
"description": "Too Many Requests (max daily resets exceeded or reset in progress)"
|
||||
},
|
||||
"503": {
|
||||
"description": "Service Unavailable (Redis reset failed; credits refunded or support needed)"
|
||||
}
|
||||
},
|
||||
"security": [{ "HTTPBearerJWT": [] }]
|
||||
}
|
||||
},
|
||||
"/api/copilot/admin/rate_limit": {
|
||||
"get": {
|
||||
"tags": ["v2", "admin", "copilot", "admin"],
|
||||
@@ -8368,7 +8402,13 @@
|
||||
"CoPilotUsageStatus": {
|
||||
"properties": {
|
||||
"daily": { "$ref": "#/components/schemas/UsageWindow" },
|
||||
"weekly": { "$ref": "#/components/schemas/UsageWindow" }
|
||||
"weekly": { "$ref": "#/components/schemas/UsageWindow" },
|
||||
"reset_cost": {
|
||||
"type": "integer",
|
||||
"title": "Reset Cost",
|
||||
"description": "Credit cost (in cents) to reset the daily limit. 0 = feature disabled.",
|
||||
"default": 0
|
||||
}
|
||||
},
|
||||
"type": "object",
|
||||
"required": ["daily", "weekly"],
|
||||
@@ -11677,6 +11717,34 @@
|
||||
"required": ["providers", "pagination"],
|
||||
"title": "ProviderResponse"
|
||||
},
|
||||
"RateLimitResetResponse": {
|
||||
"properties": {
|
||||
"success": { "type": "boolean", "title": "Success" },
|
||||
"credits_charged": {
|
||||
"type": "integer",
|
||||
"title": "Credits Charged",
|
||||
"description": "Credits charged (in cents)"
|
||||
},
|
||||
"remaining_balance": {
|
||||
"type": "integer",
|
||||
"title": "Remaining Balance",
|
||||
"description": "Credit balance after charge (in cents)"
|
||||
},
|
||||
"usage": {
|
||||
"$ref": "#/components/schemas/CoPilotUsageStatus",
|
||||
"description": "Updated usage status after reset"
|
||||
}
|
||||
},
|
||||
"type": "object",
|
||||
"required": [
|
||||
"success",
|
||||
"credits_charged",
|
||||
"remaining_balance",
|
||||
"usage"
|
||||
],
|
||||
"title": "RateLimitResetResponse",
|
||||
"description": "Response from resetting the daily rate limit."
|
||||
},
|
||||
"RecentExecution": {
|
||||
"properties": {
|
||||
"status": { "type": "string", "title": "Status" },
|
||||
|
||||
Reference in New Issue
Block a user