fix: restore weekly quotas and reconcile live health
This commit is contained in:
parent
49821451a5
commit
17c266339f
9 changed files with 278 additions and 27 deletions
|
|
@ -265,12 +265,14 @@ class QuotaSnapshot:
|
|||
return f"Обновлено: {hrs} ч назад"
|
||||
|
||||
def get_bucket_for_model(self, model_or_family: str) -> Optional[QuotaBucket]:
|
||||
"""Find the relevant quota bucket for a given model or model family."""
|
||||
"""Find the most constraining quota bucket for a model family."""
|
||||
target = model_or_family.lower()
|
||||
# Direct family match
|
||||
for b in self.buckets:
|
||||
if b.model_family and b.model_family.lower() in target:
|
||||
return b
|
||||
matches = [b for b in self.buckets if b.model_family and b.model_family.lower() in target]
|
||||
if matches:
|
||||
measured = [b for b in matches if b.remaining_percent is not None]
|
||||
if measured:
|
||||
return min(measured, key=lambda bucket: float(bucket.remaining_percent or 0.0))
|
||||
return matches[0]
|
||||
# Fallback to first available bucket
|
||||
return self.buckets[0] if self.buckets else None
|
||||
|
||||
|
|
|
|||
|
|
@ -273,6 +273,36 @@ class HealthTracker:
|
|||
|
||||
self._save_state()
|
||||
|
||||
def reconcile_measured_quota(self, profile_id: str, remaining_by_family: Dict[str, float]) -> bool:
|
||||
"""Clear stale quota-exhausted flags when the provider reports live capacity.
|
||||
|
||||
A successful quota read is authoritative for quota exhaustion, but it
|
||||
must not erase unrelated authentication or runtime failures.
|
||||
"""
|
||||
measured = {family: float(value) for family, value in remaining_by_family.items()}
|
||||
if not measured:
|
||||
return False
|
||||
with self._lock:
|
||||
record = self.get_or_create(profile_id)
|
||||
changed = False
|
||||
for family, remaining in measured.items():
|
||||
family_record = record.families.get(family)
|
||||
if remaining > 0 and family_record and family_record.state == QUOTA_EXHAUSTED:
|
||||
family_record.state = HEALTHY
|
||||
family_record.reset_at = None
|
||||
family_record.reason = None
|
||||
family_record.last_error = None
|
||||
family_record.simulated = False
|
||||
changed = True
|
||||
if all(value > 0 for value in measured.values()) and record.overall_state == QUOTA_EXHAUSTED:
|
||||
record.overall_state = HEALTHY
|
||||
record.last_error = None
|
||||
record.simulated = False
|
||||
changed = True
|
||||
if changed:
|
||||
self._save_state()
|
||||
return changed
|
||||
|
||||
def mark_quota_exhausted(
|
||||
self,
|
||||
profile_id: str,
|
||||
|
|
|
|||
|
|
@ -376,6 +376,8 @@ class HermesHubApp(ctk.CTk):
|
|||
text_color=Theme.STATUS_HEALTHY,
|
||||
)
|
||||
self.status_left.pack(side="left", padx=Theme.SPACE_LG)
|
||||
self.status_left.configure(cursor="hand2")
|
||||
self.status_left.bind("<Button-1>", lambda _event: self._show_view("health"), add="+")
|
||||
|
||||
self.global_search = ctk.CTkEntry(
|
||||
self.statusbar,
|
||||
|
|
@ -580,7 +582,7 @@ class HermesHubApp(ctk.CTk):
|
|||
|
||||
freshness = "⚠ Данные устарели" if snap.is_stale else f"Snapshot #{snap.seq}"
|
||||
self.status_left.configure(
|
||||
text=f"● {readiness.title_ru}",
|
||||
text=f"● {readiness.title_ru}{' · Подробнее' if readiness.state != 'healthy' else ''}",
|
||||
text_color=Theme.STATUS_HEALTHY
|
||||
if readiness.state == "healthy"
|
||||
else Theme.STATUS_WARNING
|
||||
|
|
@ -1108,7 +1110,7 @@ class HermesHubApp(ctk.CTk):
|
|||
|
||||
readiness = HubStateStore.get().get_snapshot().readiness
|
||||
self.status_left.configure(
|
||||
text=f"● {readiness.title_ru}",
|
||||
text=f"● {readiness.title_ru}{' · Подробнее' if readiness.state != 'healthy' else ''}",
|
||||
text_color=Theme.STATUS_HEALTHY
|
||||
if readiness.state == "healthy"
|
||||
else Theme.STATUS_WARNING
|
||||
|
|
|
|||
|
|
@ -159,6 +159,22 @@ class AccountQuotaService:
|
|||
with self._cache_lock:
|
||||
self._snapshots[key] = snap
|
||||
|
||||
if snap.source == "provider_api":
|
||||
measured_by_family: dict[str, float] = {}
|
||||
for bucket in snap.buckets:
|
||||
family = bucket.model_family
|
||||
remaining = bucket.remaining_percent
|
||||
if not family or remaining is None:
|
||||
continue
|
||||
measured_by_family[family] = min(measured_by_family.get(family, 100.0), float(remaining))
|
||||
if measured_by_family:
|
||||
try:
|
||||
from .router_engine import get_router_engine
|
||||
|
||||
get_router_engine().health.reconcile_measured_quota(profile_id, measured_by_family)
|
||||
except Exception as exc:
|
||||
logger.debug("Could not reconcile live quota health for %s: %s", profile_id, exc)
|
||||
|
||||
# Notify listeners
|
||||
for listener in list(self._listeners):
|
||||
try:
|
||||
|
|
@ -362,9 +378,9 @@ class AccountQuotaService:
|
|||
auth_data["project_id"] = project_id
|
||||
ProfileAuthManager.save_profile_auth("antigravity", profile_id, auth_data)
|
||||
|
||||
def _fetch(token: str) -> dict[str, Any]:
|
||||
def _fetch(token: str, operation: str) -> dict[str, Any]:
|
||||
request = urllib.request.Request(
|
||||
"https://daily-cloudcode-pa.googleapis.com/v1internal:fetchAvailableModels",
|
||||
f"https://daily-cloudcode-pa.googleapis.com/v1internal:{operation}",
|
||||
data=json.dumps({"project": project_id}).encode("utf-8"),
|
||||
headers={
|
||||
"Authorization": f"Bearer {token}",
|
||||
|
|
@ -376,12 +392,67 @@ class AccountQuotaService:
|
|||
with urllib.request.urlopen(request, timeout=30) as response:
|
||||
return json.loads(response.read().decode("utf-8") or "{}")
|
||||
|
||||
def _fetch_with_refresh(operation: str) -> dict[str, Any]:
|
||||
nonlocal access_token
|
||||
try:
|
||||
return _fetch(str(access_token), operation)
|
||||
except urllib.error.HTTPError as exc:
|
||||
if exc.code != 401 or not refresh_token:
|
||||
raise
|
||||
access_token = _refresh_and_save()
|
||||
return _fetch(access_token, operation)
|
||||
|
||||
# This is the endpoint used by the Antigravity usage screen. It
|
||||
# exposes the four semantic buckets: Gemini and Claude/GPT, each with
|
||||
# a five-hour and weekly window. Treat it as best-effort so older
|
||||
# accounts can still fall back to per-model capacity below.
|
||||
try:
|
||||
payload = _fetch(str(access_token))
|
||||
except urllib.error.HTTPError as exc:
|
||||
if exc.code != 401 or not refresh_token:
|
||||
raise
|
||||
payload = _fetch(_refresh_and_save())
|
||||
summary_payload = _fetch_with_refresh("retrieveUserQuotaSummary")
|
||||
except Exception as exc:
|
||||
logger.info("Grouped Antigravity quota unavailable for %s: %s", profile_id, exc)
|
||||
summary_payload = {}
|
||||
|
||||
summary_groups = summary_payload.get("groups") if isinstance(summary_payload, dict) else None
|
||||
grouped: dict[tuple[str, str], QuotaBucket] = {}
|
||||
if isinstance(summary_groups, list):
|
||||
for group in summary_groups:
|
||||
if not isinstance(group, dict):
|
||||
continue
|
||||
group_name = str(group.get("displayName") or group.get("description") or "").lower()
|
||||
family = "gemini" if "gemini" in group_name else "claude" if "claude" in group_name else None
|
||||
family_label = "Gemini" if family == "gemini" else "Claude/GPT"
|
||||
if family is None:
|
||||
continue
|
||||
for bucket_data in group.get("buckets") or []:
|
||||
if not isinstance(bucket_data, dict):
|
||||
continue
|
||||
remaining_fraction = bucket_data.get("remainingFraction")
|
||||
window_raw = str(bucket_data.get("window") or "").lower()
|
||||
window = "7d" if "week" in window_raw or window_raw == "7d" else "5h" if "5" in window_raw else ""
|
||||
if not isinstance(remaining_fraction, (int, float)) or not window:
|
||||
continue
|
||||
window_label = "неделя" if window == "7d" else "5 часов"
|
||||
grouped[(family, window)] = QuotaBucket(
|
||||
id=f"antigravity.{family}.{window}",
|
||||
display_name=f"{family_label} • {window_label}",
|
||||
model_family=family,
|
||||
remaining_percent=max(0.0, min(100.0, float(remaining_fraction) * 100.0)),
|
||||
reset_at=_parse_datetime(bucket_data.get("resetTime")),
|
||||
period=window,
|
||||
unit="model capacity",
|
||||
scope="model_family",
|
||||
)
|
||||
ordered_keys = (("claude", "5h"), ("gemini", "5h"), ("claude", "7d"), ("gemini", "7d"))
|
||||
if all(key in grouped for key in ordered_keys):
|
||||
return QuotaSnapshot(
|
||||
account_id=profile_id,
|
||||
provider="antigravity",
|
||||
buckets=[grouped[key] for key in ordered_keys],
|
||||
fetched_at=now,
|
||||
source="provider_api",
|
||||
)
|
||||
|
||||
payload = _fetch_with_refresh("fetchAvailableModels")
|
||||
|
||||
models = payload.get("models") or {}
|
||||
if not isinstance(models, dict):
|
||||
|
|
|
|||
|
|
@ -518,6 +518,16 @@ class DashboardView(ctk.CTkFrame):
|
|||
return f"{bucket.remaining_percent:.0f}%", float(bucket.remaining_percent)
|
||||
return "Нет данных API", None
|
||||
|
||||
@staticmethod
|
||||
def _agent_quota_measurement(snapshot: HubSnapshot, agent: Any) -> tuple[str, Optional[float]]:
|
||||
quota = snapshot.quotas.get(agent.assigned_profile_id)
|
||||
if quota and not getattr(quota, "is_estimated", True):
|
||||
bucket = quota.get_bucket_for_model(agent.model)
|
||||
if bucket and bucket.remaining_percent is not None:
|
||||
remaining = float(bucket.remaining_percent)
|
||||
return bucket.formatted_remaining(), remaining
|
||||
return agent.active_quota_label or "Нет данных API", None
|
||||
|
||||
@staticmethod
|
||||
def _sync_endpoint_cards(
|
||||
slots: list[Any],
|
||||
|
|
@ -626,22 +636,25 @@ class DashboardView(ctk.CTkFrame):
|
|||
self.route_diagram.orchestrator.update_node("Не назначен", "Н/Д", False)
|
||||
|
||||
agents = [agent for agent in snapshot.agents if not agent.is_main_orchestrator][:5]
|
||||
self._sync_endpoint_cards(
|
||||
self.route_diagram.agent_slots,
|
||||
self._agent_cards,
|
||||
(
|
||||
agent_items = []
|
||||
for agent in agents:
|
||||
agent_quota_text, agent_quota_percent = self._agent_quota_measurement(snapshot, agent)
|
||||
agent_items.append(
|
||||
(
|
||||
agent.role_id,
|
||||
agent.provider,
|
||||
agent.role_name_ru,
|
||||
f"{agent.provider_display_name} • {agent.model}",
|
||||
"Здорово" if agent.is_active else agent.status_label_ru,
|
||||
agent.active_quota_label or "Нет данных API",
|
||||
None,
|
||||
agent_quota_text,
|
||||
agent_quota_percent,
|
||||
"healthy" if agent.is_active else "warning",
|
||||
)
|
||||
for agent in agents
|
||||
),
|
||||
)
|
||||
self._sync_endpoint_cards(
|
||||
self.route_diagram.agent_slots,
|
||||
self._agent_cards,
|
||||
agent_items,
|
||||
)
|
||||
for agent in agents:
|
||||
card = self._agent_cards.get(agent.role_id)
|
||||
|
|
|
|||
|
|
@ -511,10 +511,17 @@ class UnifiedHealthService:
|
|||
dead_roles = 0
|
||||
warnings: List[str] = []
|
||||
|
||||
total_accounts = sum(len(profs) for profs in profiles_by_prov.values())
|
||||
connected_accounts = sum(
|
||||
1 for profs in profiles_by_prov.values() for p in profs if p.auth_state == "AUTHENTICATED"
|
||||
)
|
||||
# Empty/cold placeholders are capacity for future accounts, not broken
|
||||
# accounts. They remain visible in provider slot counts but must not
|
||||
# downgrade a system whose configured routes are all operational.
|
||||
configured_profiles = [
|
||||
profile
|
||||
for profiles in profiles_by_prov.values()
|
||||
for profile in profiles
|
||||
if not profile.is_empty_slot
|
||||
]
|
||||
total_accounts = len(configured_profiles)
|
||||
connected_accounts = sum(1 for profile in configured_profiles if profile.auth_state == "AUTHENTICATED")
|
||||
|
||||
providers_online = sum(
|
||||
1 for profs in profiles_by_prov.values() if any(p.health_state == STATUS_HEALTHY for p in profs)
|
||||
|
|
|
|||
|
|
@ -224,6 +224,69 @@ def test_antigravity_separate_claude_and_gemini_buckets():
|
|||
assert b_g.remaining_percent == 87.0
|
||||
|
||||
|
||||
def test_antigravity_grouped_summary_includes_five_hour_and_weekly_buckets():
|
||||
service = AccountQuotaService()
|
||||
response = MagicMock()
|
||||
response.__enter__.return_value.read.return_value = json.dumps(
|
||||
{
|
||||
"groups": [
|
||||
{
|
||||
"displayName": "Gemini Models",
|
||||
"buckets": [
|
||||
{
|
||||
"bucketId": "gemini-weekly",
|
||||
"window": "weekly",
|
||||
"remainingFraction": 0.73,
|
||||
"resetTime": "2026-08-27T00:00:00Z",
|
||||
},
|
||||
{
|
||||
"bucketId": "gemini-5h",
|
||||
"window": "5h",
|
||||
"remainingFraction": 0.91,
|
||||
"resetTime": "2026-08-23T05:00:00Z",
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
"displayName": "Claude and GPT models",
|
||||
"buckets": [
|
||||
{
|
||||
"bucketId": "3p-weekly",
|
||||
"window": "weekly",
|
||||
"remainingFraction": 0.44,
|
||||
"resetTime": "2026-08-27T00:00:00Z",
|
||||
},
|
||||
{
|
||||
"bucketId": "3p-5h",
|
||||
"window": "5h",
|
||||
"remainingFraction": 0.82,
|
||||
"resetTime": "2026-08-23T05:00:00Z",
|
||||
},
|
||||
],
|
||||
},
|
||||
]
|
||||
}
|
||||
).encode()
|
||||
|
||||
with patch(
|
||||
"antigravity_provider.router.quota_collector.urllib.request.urlopen", return_value=response
|
||||
) as urlopen:
|
||||
snapshot = service._collect_antigravity_quota(
|
||||
"ag-w1", {"token": {"access_token": "access-token"}, "project_id": "project-1"}
|
||||
)
|
||||
|
||||
assert urlopen.call_count == 1
|
||||
assert [bucket.id for bucket in snapshot.buckets] == [
|
||||
"antigravity.claude.5h",
|
||||
"antigravity.gemini.5h",
|
||||
"antigravity.claude.7d",
|
||||
"antigravity.gemini.7d",
|
||||
]
|
||||
assert [bucket.remaining_percent for bucket in snapshot.buckets] == [82.0, 91.0, 44.0, 73.0]
|
||||
assert snapshot.get_bucket_for_model("gemini-3.1-pro").remaining_percent == 73.0
|
||||
assert snapshot.get_bucket_for_model("claude-sonnet-4-6").remaining_percent == 44.0
|
||||
|
||||
|
||||
def test_antigravity_refreshes_expired_token_before_project_discovery():
|
||||
service = AccountQuotaService()
|
||||
response = MagicMock()
|
||||
|
|
@ -306,6 +369,17 @@ def test_health_tracker_antigravity_claude_exhaustion_does_not_block_gemini(tmp_
|
|||
|
||||
# Gemini should remain healthy!
|
||||
assert tracker.is_healthy("ag-w1", "gemini-2.5-pro") is True
|
||||
|
||||
|
||||
def test_live_measured_quota_clears_stale_exhaustion(tmp_path):
|
||||
tracker = HealthTracker(state_file=tmp_path / "router_state.json")
|
||||
tracker.mark_quota_exhausted("ag-w1", "gemini-3.1-pro", duration=3600, reason="old 429")
|
||||
assert tracker.is_healthy("ag-w1", "gemini-3.1-pro") is False
|
||||
|
||||
changed = tracker.reconcile_measured_quota("ag-w1", {"gemini": 73.0, "claude": 44.0})
|
||||
|
||||
assert changed is True
|
||||
assert tracker.is_healthy("ag-w1", "gemini-3.1-pro") is True
|
||||
assert tracker.is_healthy("ag-w1", "gemini-2.5-flash") is True
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -212,6 +212,58 @@ def test_agent_quota_and_failover_reason_are_bound_to_their_models(ui_root) -> N
|
|||
team_card.destroy()
|
||||
|
||||
|
||||
def test_dashboard_agent_quota_measurement_drives_progress_percent() -> None:
|
||||
agent = AgentViewModel(
|
||||
role_id="coder-primary",
|
||||
role_name_ru="Кодер 1",
|
||||
role_description_ru="Основной кодер",
|
||||
assigned_profile_id="ag-w1",
|
||||
assigned_display_name="Primary",
|
||||
provider="antigravity",
|
||||
provider_display_name="Google Antigravity",
|
||||
model="gemini-3.1-pro",
|
||||
account_identity="user@example.test",
|
||||
routing_position="Primary",
|
||||
status="healthy",
|
||||
status_label_ru="Работает",
|
||||
is_active=True,
|
||||
is_main_orchestrator=False,
|
||||
active_quota_status="healthy",
|
||||
active_quota_label="Осталось 73%",
|
||||
)
|
||||
snapshot = HubSnapshot(
|
||||
generation=1,
|
||||
seq=1,
|
||||
timestamp=time.time(),
|
||||
profiles_by_provider={},
|
||||
all_profiles={},
|
||||
readiness=_readiness(),
|
||||
agents=[agent],
|
||||
providers=[],
|
||||
routing={},
|
||||
quotas={
|
||||
"ag-w1": QuotaSnapshot(
|
||||
account_id="ag-w1",
|
||||
provider="antigravity",
|
||||
source="provider_api",
|
||||
buckets=[
|
||||
QuotaBucket(
|
||||
id="antigravity.gemini.7d",
|
||||
display_name="Gemini • неделя",
|
||||
model_family="gemini",
|
||||
remaining_percent=73.0,
|
||||
)
|
||||
],
|
||||
)
|
||||
},
|
||||
)
|
||||
|
||||
label, percent = DashboardView._agent_quota_measurement(snapshot, agent)
|
||||
|
||||
assert label == "Осталось 73%"
|
||||
assert percent == 73.0
|
||||
|
||||
|
||||
@pytest.mark.ui
|
||||
def test_stale_snapshot_is_visibly_marked_with_sequence(ui_root) -> None:
|
||||
snapshot = HubSnapshot(
|
||||
|
|
|
|||
|
|
@ -109,7 +109,7 @@ def test_system_readiness_calculation():
|
|||
assert isinstance(readiness, SystemReadiness)
|
||||
assert readiness.state in (READINESS_HEALTHY, READINESS_LIMITED, READINESS_DEGRADED, READINESS_CRITICAL)
|
||||
assert readiness.total_roles > 0
|
||||
assert readiness.total_accounts >= 16
|
||||
assert readiness.total_accounts >= readiness.accounts_connected_count
|
||||
assert readiness.title_ru
|
||||
assert readiness.summary_ru
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue