350 lines
13 KiB
Python
350 lines
13 KiB
Python
"""Hermes Hub — Unified HubStateStore & Immutable HubSnapshot Layer.
|
|
|
|
Provides normalized state management, single-scan snapshot generation,
|
|
request deduplication, generation tracking, and delta event publishing.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import threading
|
|
import time
|
|
from dataclasses import dataclass, field, replace
|
|
from typing import Any, Dict, List, Optional, Tuple
|
|
|
|
from antigravity_provider.router.event_bus import (
|
|
EventBus,
|
|
EVENT_ACCOUNT_UPDATED,
|
|
EVENT_ACCOUNT_ADDED,
|
|
EVENT_ACCOUNT_REMOVED,
|
|
EVENT_ACCOUNT_AUTH_CHANGED,
|
|
EVENT_QUOTA_UPDATED,
|
|
EVENT_ROUTING_UPDATED,
|
|
EVENT_SYSTEM_READINESS_CHANGED,
|
|
EVENT_REFRESH_STARTED,
|
|
EVENT_REFRESH_COMPLETED,
|
|
EVENT_REFRESH_FAILED,
|
|
)
|
|
from antigravity_provider.router.unified_health import (
|
|
UnifiedHealthService,
|
|
ProfileViewModel,
|
|
SystemReadiness,
|
|
AgentViewModel,
|
|
ProviderSummary,
|
|
RolePipeline,
|
|
)
|
|
from antigravity_provider.router.quota_collector import AccountQuotaService
|
|
from antigravity_provider.router.router_config import load_router_config
|
|
from antigravity_provider.router.auto_assigner import AutoAssigner
|
|
|
|
logger = logging.getLogger("hermes.router.state_store")
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class HubSnapshot:
|
|
"""Immutable normalized snapshot of the entire Hermes Hub state at a specific generation."""
|
|
generation: int
|
|
seq: int
|
|
timestamp: float
|
|
profiles_by_provider: Dict[str, List[ProfileViewModel]]
|
|
all_profiles: Dict[str, ProfileViewModel]
|
|
readiness: SystemReadiness
|
|
agents: List[AgentViewModel]
|
|
providers: List[ProviderSummary]
|
|
routing: Dict[str, RolePipeline]
|
|
quotas: Dict[str, Any]
|
|
metrics: Dict[str, Any] = field(default_factory=dict)
|
|
is_stale: bool = False
|
|
|
|
def get_profile(self, profile_id: str) -> Optional[ProfileViewModel]:
|
|
return self.all_profiles.get(profile_id)
|
|
|
|
def get_provider_profiles(self, provider: str) -> List[ProfileViewModel]:
|
|
return list(self.profiles_by_provider.get(provider, []))
|
|
|
|
def get_role_pipeline(self, role_id: str) -> Optional[RolePipeline]:
|
|
return self.routing.get(role_id)
|
|
|
|
|
|
class HubStateStore:
|
|
"""Thread-safe central state store managing the canonical HubSnapshot and delta updates."""
|
|
|
|
_instance: Optional[HubStateStore] = None
|
|
_instance_lock = threading.Lock()
|
|
|
|
def __init__(self) -> None:
|
|
self._lock = threading.RLock()
|
|
self._generation: int = 0
|
|
self._current_snapshot: Optional[HubSnapshot] = None
|
|
self._pending_refreshes: Dict[str, float] = {}
|
|
self._latest_applied_seq: int = 0
|
|
self._seq_counter: int = 0
|
|
|
|
# Observability counters
|
|
self.refresh_runs_total: int = 0
|
|
self.refresh_skipped_total: int = 0
|
|
self.refresh_deduplicated_total: int = 0
|
|
self.refresh_failures_total: int = 0
|
|
self.account_updates_total: int = 0
|
|
self.quota_updates_total: int = 0
|
|
|
|
@classmethod
|
|
def get(cls) -> HubStateStore:
|
|
if cls._instance is None:
|
|
with cls._instance_lock:
|
|
if cls._instance is None:
|
|
cls._instance = cls()
|
|
return cls._instance
|
|
|
|
def next_seq(self) -> int:
|
|
with self._lock:
|
|
self._seq_counter += 1
|
|
return self._seq_counter
|
|
|
|
def get_snapshot(self) -> HubSnapshot:
|
|
"""Return the current cached snapshot. Generates an initial snapshot if none exists."""
|
|
with self._lock:
|
|
if self._current_snapshot is not None:
|
|
return self._current_snapshot
|
|
return self.refresh(force_scan=False)
|
|
|
|
def refresh(self, force_scan: bool = True, seq: Optional[int] = None) -> HubSnapshot:
|
|
"""Build outside the store lock, then atomically apply only the newest result."""
|
|
request_seq = seq if seq is not None else self.next_seq()
|
|
t0 = time.time()
|
|
|
|
# Slow disk/provider reads must not hold the store lock. More recent
|
|
# requests may complete while this build is running.
|
|
uh_service = UnifiedHealthService.get()
|
|
profiles_by_prov = uh_service.scan_all(force=force_scan)
|
|
all_profs = {
|
|
profile.profile_id: profile
|
|
for profiles in profiles_by_prov.values()
|
|
for profile in profiles
|
|
}
|
|
readiness = uh_service.get_system_readiness()
|
|
agents = uh_service.get_agent_view_models()
|
|
providers = uh_service.get_provider_summaries()
|
|
routing = uh_service.get_routing_pipelines()
|
|
quota_service = AccountQuotaService.get()
|
|
quotas_map = {
|
|
profile_id: quota_service.get_snapshot(profile.provider, profile_id)
|
|
for profile_id, profile in all_profs.items()
|
|
if profile.auth_state == "AUTHENTICATED"
|
|
}
|
|
|
|
with self._lock:
|
|
self.refresh_runs_total += 1
|
|
if request_seq < self._latest_applied_seq:
|
|
logger.info(
|
|
"Discarding late refresh result (seq %d < applied %d)",
|
|
request_seq,
|
|
self._latest_applied_seq,
|
|
)
|
|
self.refresh_skipped_total += 1
|
|
return self._current_snapshot or self._build_empty_snapshot()
|
|
|
|
self._latest_applied_seq = request_seq
|
|
self._generation += 1
|
|
gen = self._generation
|
|
metrics = {
|
|
"generation": gen,
|
|
"seq": request_seq,
|
|
"duration_ms": round((time.time() - t0) * 1000, 2),
|
|
"total_profiles": len(all_profs),
|
|
"authenticated_profiles": sum(
|
|
1 for profile in all_profs.values() if profile.auth_state == "AUTHENTICATED"
|
|
),
|
|
"refresh_runs_total": self.refresh_runs_total,
|
|
"refresh_deduplicated_total": self.refresh_deduplicated_total,
|
|
}
|
|
snapshot = HubSnapshot(
|
|
generation=gen,
|
|
seq=request_seq,
|
|
timestamp=time.time(),
|
|
profiles_by_provider=profiles_by_prov,
|
|
all_profiles=all_profs,
|
|
readiness=readiness,
|
|
agents=agents,
|
|
providers=providers,
|
|
routing=routing,
|
|
quotas=quotas_map,
|
|
metrics=metrics,
|
|
is_stale=False,
|
|
)
|
|
self._current_snapshot = snapshot
|
|
|
|
# Emit snapshot update on EventBus
|
|
EventBus.get().publish(EVENT_SYSTEM_READINESS_CHANGED, readiness)
|
|
EventBus.get().publish(
|
|
EVENT_REFRESH_COMPLETED,
|
|
{"generation": gen, "seq": request_seq, "duration_ms": metrics["duration_ms"]},
|
|
)
|
|
return snapshot
|
|
|
|
def _build_empty_snapshot(self) -> HubSnapshot:
|
|
return HubSnapshot(
|
|
generation=0,
|
|
seq=0,
|
|
timestamp=time.time(),
|
|
profiles_by_provider={},
|
|
all_profiles={},
|
|
readiness=SystemReadiness(state="limited", title_ru="Инициализация", description_ru="", accounts_connected_count=0, total_accounts=0, roles_ready_count=0, total_roles=0, providers_ready_count=0, total_providers=0),
|
|
agents=[],
|
|
providers=[],
|
|
routing={},
|
|
quotas={},
|
|
metrics={},
|
|
is_stale=True,
|
|
)
|
|
|
|
def _apply_profile_delta(self, profile: ProfileViewModel) -> HubSnapshot:
|
|
"""Copy-on-write replacement of exactly one profile in the snapshot."""
|
|
with self._lock:
|
|
current = self._current_snapshot or self._build_empty_snapshot()
|
|
all_profiles = dict(current.all_profiles)
|
|
all_profiles[profile.profile_id] = profile
|
|
grouped = {provider: list(items) for provider, items in current.profiles_by_provider.items()}
|
|
provider_profiles = grouped.setdefault(profile.provider, [])
|
|
for index, existing in enumerate(provider_profiles):
|
|
if existing.profile_id == profile.profile_id:
|
|
provider_profiles[index] = profile
|
|
break
|
|
else:
|
|
provider_profiles.append(profile)
|
|
self._generation += 1
|
|
seq = self.next_seq()
|
|
self._latest_applied_seq = seq
|
|
updated = replace(
|
|
current,
|
|
generation=self._generation,
|
|
seq=seq,
|
|
timestamp=time.time(),
|
|
profiles_by_provider=grouped,
|
|
all_profiles=all_profiles,
|
|
)
|
|
self._current_snapshot = updated
|
|
return updated
|
|
|
|
def apply_delta_account_updated(
|
|
self,
|
|
profile_id: str,
|
|
profile: Optional[ProfileViewModel] = None,
|
|
) -> None:
|
|
"""Update one account and publish a profile-keyed event without a global scan."""
|
|
self.account_updates_total += 1
|
|
if profile is None:
|
|
cached = UnifiedHealthService.get().get_cached_profiles()
|
|
profile = next(
|
|
(item for items in cached.values() for item in items if item.profile_id == profile_id),
|
|
None,
|
|
)
|
|
if profile is None:
|
|
logger.warning("Cannot apply account delta for unknown profile %s", profile_id)
|
|
return
|
|
snapshot = self._apply_profile_delta(profile)
|
|
EventBus.get().publish(
|
|
EVENT_ACCOUNT_UPDATED,
|
|
{
|
|
"provider": profile.provider,
|
|
"profile_id": profile_id,
|
|
"profile": profile,
|
|
"generation": snapshot.generation,
|
|
"seq": snapshot.seq,
|
|
},
|
|
)
|
|
|
|
def apply_delta_account_added(self, profile: ProfileViewModel) -> None:
|
|
snapshot = self._apply_profile_delta(profile)
|
|
EventBus.get().publish(
|
|
EVENT_ACCOUNT_ADDED,
|
|
{
|
|
"provider": profile.provider,
|
|
"profile_id": profile.profile_id,
|
|
"profile": profile,
|
|
"generation": snapshot.generation,
|
|
"seq": snapshot.seq,
|
|
},
|
|
)
|
|
|
|
def apply_delta_account_removed(self, provider: str, profile_id: str) -> None:
|
|
with self._lock:
|
|
current = self._current_snapshot or self._build_empty_snapshot()
|
|
all_profiles = dict(current.all_profiles)
|
|
all_profiles.pop(profile_id, None)
|
|
grouped = {key: list(value) for key, value in current.profiles_by_provider.items()}
|
|
grouped[provider] = [item for item in grouped.get(provider, []) if item.profile_id != profile_id]
|
|
quotas = dict(current.quotas)
|
|
quotas.pop(profile_id, None)
|
|
self._generation += 1
|
|
seq = self.next_seq()
|
|
self._latest_applied_seq = seq
|
|
updated = replace(
|
|
current,
|
|
generation=self._generation,
|
|
seq=seq,
|
|
timestamp=time.time(),
|
|
profiles_by_provider=grouped,
|
|
all_profiles=all_profiles,
|
|
quotas=quotas,
|
|
)
|
|
self._current_snapshot = updated
|
|
EventBus.get().publish(
|
|
EVENT_ACCOUNT_REMOVED,
|
|
{"provider": provider, "profile_id": profile_id, "generation": updated.generation, "seq": seq},
|
|
)
|
|
|
|
def publish_auth_changed(self, profile: ProfileViewModel) -> None:
|
|
snapshot = self._apply_profile_delta(profile)
|
|
EventBus.get().publish(
|
|
EVENT_ACCOUNT_AUTH_CHANGED,
|
|
{
|
|
"provider": profile.provider,
|
|
"profile_id": profile.profile_id,
|
|
"auth_state": profile.auth_state,
|
|
"profile": profile,
|
|
"generation": snapshot.generation,
|
|
"seq": snapshot.seq,
|
|
},
|
|
)
|
|
|
|
def apply_delta_quota_updated(self, provider: str, profile_id: str, quota_snap: Any) -> None:
|
|
"""Apply instant runtime quota change (e.g. 429 received during inference)."""
|
|
with self._lock:
|
|
self.quota_updates_total += 1
|
|
current = self._current_snapshot or self._build_empty_snapshot()
|
|
quotas = dict(current.quotas)
|
|
quotas[profile_id] = quota_snap
|
|
all_profiles = dict(current.all_profiles)
|
|
profile = all_profiles.get(profile_id)
|
|
grouped = {provider_key: list(items) for provider_key, items in current.profiles_by_provider.items()}
|
|
if profile is not None:
|
|
updated_profile = replace(profile, quota_snapshot=quota_snap)
|
|
all_profiles[profile_id] = updated_profile
|
|
grouped[profile.provider] = [
|
|
updated_profile if item.profile_id == profile_id else item
|
|
for item in grouped.get(profile.provider, [])
|
|
]
|
|
self._generation += 1
|
|
seq = self.next_seq()
|
|
self._latest_applied_seq = seq
|
|
updated = replace(
|
|
current,
|
|
generation=self._generation,
|
|
seq=seq,
|
|
timestamp=time.time(),
|
|
profiles_by_provider=grouped,
|
|
quotas=quotas,
|
|
all_profiles=all_profiles,
|
|
)
|
|
self._current_snapshot = updated
|
|
|
|
EventBus.get().publish(
|
|
EVENT_QUOTA_UPDATED,
|
|
{
|
|
"provider": provider,
|
|
"profile_id": profile_id,
|
|
"quota_snapshot": quota_snap,
|
|
"generation": updated.generation,
|
|
"seq": updated.seq,
|
|
},
|
|
)
|