diff --git a/src/antigravity_provider/__init__.py b/src/antigravity_provider/__init__.py new file mode 100644 index 0000000..dca5cf1 --- /dev/null +++ b/src/antigravity_provider/__init__.py @@ -0,0 +1,15 @@ +"""Hermes Hub provider package. + +This explicit package boundary prevents Python from merging the repository and +an older installed plugin copy as a namespace package. A process therefore +loads one coherent source tree instead of a mixture of versions. +""" +from __future__ import annotations + +from pathlib import Path + +from .version import __version__ + +PACKAGE_ROOT = Path(__file__).resolve().parent + +__all__ = ["PACKAGE_ROOT", "__version__"] diff --git a/src/antigravity_provider/router/__init__.py b/src/antigravity_provider/router/__init__.py index e4e6d2d..33f1061 100644 --- a/src/antigravity_provider/router/__init__.py +++ b/src/antigravity_provider/router/__init__.py @@ -16,11 +16,6 @@ from .health_tracker import ( ) from .session_affinity import LeaseManager, SessionAffinityRecord, SessionAffinityTracker from .router_engine import RouterEngine, get_router_engine -from .capability.capability_matrix import CapabilityMatrix, ModelCapability -from .supervisor.lifecycle_supervisor import LifecycleSupervisor -from .skills.skill_registry import UnifiedSkill, UnifiedSkillRegistry - -SkillRegistry = UnifiedSkillRegistry __all__ = [ "RouterConfig", @@ -42,10 +37,4 @@ __all__ = [ "LeaseManager", "RouterEngine", "get_router_engine", - "CapabilityMatrix", - "ModelCapability", - "LifecycleSupervisor", - "SkillRegistry", - "UnifiedSkillRegistry", - "UnifiedSkill", ] diff --git a/src/antigravity_provider/router/account_identity.py b/src/antigravity_provider/router/account_identity.py index 62a8bb1..51ff8a9 100644 --- a/src/antigravity_provider/router/account_identity.py +++ b/src/antigravity_provider/router/account_identity.py @@ -148,7 +148,21 @@ class QuotaBucket: reset_at: Optional[datetime] = None reset_in_seconds: Optional[int] = None period: Optional[str] = None # "5h", "7d", "30d", "sliding" - status: str = "healthy" # "healthy", "warning", "exhausted", "unknown" + unit: Optional[str] = None # requests, tokens, tasks, currency, or provider-defined + scope: Optional[str] = None # account, model_family, organization, or provider-defined + status: str = "unknown" # "healthy", "warning", "exhausted", "unknown" + + @property + def label(self) -> str: + return self.display_name + + @property + def used(self) -> Optional[int]: + return self.used_absolute + + @property + def limit(self) -> Optional[int]: + return self.limit_absolute def __post_init__(self): # Auto-reconcile percentages @@ -187,7 +201,7 @@ class QuotaBucket: return f"Осталось {self.remaining_percent:.0f}%" if self.used_percent is not None: return f"Использовано {self.used_percent:.0f}%" - return "Доступна" + return "Н/Д" def formatted_reset(self) -> Optional[str]: """User-facing reset time string.""" @@ -227,7 +241,13 @@ class QuotaSnapshot: @property def is_estimated(self) -> bool: """True if values are baseline or locally estimated rather than measured by live server API.""" - return self.source in ("baseline", "estimated", "unconfigured", "local_heuristic") + return self.source in ( + "baseline", + "estimated", + "unconfigured", + "local_heuristic", + "runtime_error", + ) def is_stale(self) -> bool: delta = _utc_now() - self.fetched_at diff --git a/src/antigravity_provider/router/capability/capability_matrix.py b/src/antigravity_provider/router/capability/capability_matrix.py deleted file mode 100644 index d95e064..0000000 --- a/src/antigravity_provider/router/capability/capability_matrix.py +++ /dev/null @@ -1,181 +0,0 @@ -"""Model Family Capability Matrix for intelligent routing and capability matching.""" -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import Any, Dict, List, Optional, Set - - -@dataclass -class ModelCapability: - model_id: str - family: str - provider: str - context_window: int - max_output_tokens: int - supports_tools: bool = True - supports_vision: bool = False - supports_json_schema: bool = True - supports_reasoning: bool = False - supports_responses_api: bool = False - relative_cost_score: int = 1 # 1 (low) to 5 (high) - speed_score: int = 4 # 1 (slow) to 5 (fast) - - -class CapabilityMatrix: - """Matrix defining capabilities of all supported model families.""" - - _instance: Optional[CapabilityMatrix] = None - - def __init__(self): - self._models: Dict[str, ModelCapability] = {} - self._populate_matrix() - - @classmethod - def get(cls) -> CapabilityMatrix: - if cls._instance is None: - cls._instance = cls() - return cls._instance - - def _populate_matrix(self): - # ── Google Antigravity ── - self.register(ModelCapability( - model_id="gemini-2.5-pro", - family="gemini-2.5-pro", - provider="antigravity", - context_window=1000000, - max_output_tokens=65536, - supports_tools=True, - supports_vision=True, - supports_json_schema=True, - supports_reasoning=True, - relative_cost_score=4, - speed_score=3, - )) - self.register(ModelCapability( - model_id="gemini-2.5-flash", - family="gemini-2.5-flash", - provider="antigravity", - context_window=1000000, - max_output_tokens=65536, - supports_tools=True, - supports_vision=True, - supports_json_schema=True, - supports_reasoning=False, - relative_cost_score=1, - speed_score=5, - )) - self.register(ModelCapability( - model_id="gemini-2.5-flash-thinking", - family="gemini-2.5-flash-thinking", - provider="antigravity", - context_window=1000000, - max_output_tokens=65536, - supports_tools=True, - supports_vision=True, - supports_json_schema=True, - supports_reasoning=True, - relative_cost_score=2, - speed_score=4, - )) - - # ── OpenAI Codex ── - self.register(ModelCapability( - model_id="gpt-5.3-codex", - family="gpt-5.3-codex", - provider="openai-codex", - context_window=200000, - max_output_tokens=32768, - supports_tools=True, - supports_vision=True, - supports_json_schema=True, - supports_reasoning=True, - supports_responses_api=True, - relative_cost_score=5, - speed_score=3, - )) - self.register(ModelCapability( - model_id="gpt-5.1-codex-mini", - family="gpt-5.1-codex-mini", - provider="openai-codex", - context_window=128000, - max_output_tokens=16384, - supports_tools=True, - supports_vision=False, - supports_json_schema=True, - supports_reasoning=False, - supports_responses_api=True, - relative_cost_score=2, - speed_score=5, - )) - - # ── OpenCode Go ── - self.register(ModelCapability( - model_id="opencode-go-3", - family="opencode-go-3", - provider="opencode-go", - context_window=64000, - max_output_tokens=8192, - supports_tools=True, - supports_vision=False, - supports_json_schema=True, - supports_reasoning=False, - relative_cost_score=1, - speed_score=5, - )) - - # ── DeepSeek (Roadmap) ── - self.register(ModelCapability( - model_id="deepseek-chat", - family="deepseek-chat", - provider="deepseek", - context_window=64000, - max_output_tokens=8192, - supports_tools=True, - supports_vision=False, - supports_json_schema=True, - supports_reasoning=False, - supports_responses_api=True, - relative_cost_score=1, - speed_score=4, - )) - self.register(ModelCapability( - model_id="deepseek-reasoner", - family="deepseek-reasoner", - provider="deepseek", - context_window=64000, - max_output_tokens=8192, - supports_tools=False, - supports_vision=False, - supports_json_schema=False, - supports_reasoning=True, - supports_responses_api=True, - relative_cost_score=1, - speed_score=2, - )) - - def register(self, cap: ModelCapability): - self._models[cap.model_id] = cap - - def get_capability(self, model_id: str) -> Optional[ModelCapability]: - return self._models.get(model_id) - - def find_best_model_for_role( - self, - required_tools: bool = True, - required_reasoning: bool = False, - min_context: int = 32000, - ) -> List[ModelCapability]: - """Filter and rank models meeting specific role constraints.""" - candidates = [] - for m in self._models.values(): - if required_tools and not m.supports_tools: - continue - if required_reasoning and not m.supports_reasoning: - continue - if m.context_window < min_context: - continue - candidates.append(m) - - # Rank by speed and cost - candidates.sort(key=lambda x: (x.supports_reasoning, -x.relative_cost_score, x.speed_score), reverse=True) - return candidates diff --git a/src/antigravity_provider/router/model_registry.py b/src/antigravity_provider/router/model_registry.py index 780b77e..db8bdf9 100644 --- a/src/antigravity_provider/router/model_registry.py +++ b/src/antigravity_provider/router/model_registry.py @@ -49,6 +49,7 @@ class RoleRequirements: latency_priority: float = 0.5 # 0.0 (ignore latency) to 1.0 (maximize speed) reasoning_priority: float = 0.5 # 0.0 to 1.0 quality_priority: float = 0.5 # 0.0 to 1.0 + quota_priority: float = 1.0 # Prefer model families with measured quota remaining diversity_priority: float = 0.0 # 0.0 to 1.0 (prefer different provider/family from reference) allow_model_fallback: bool = True @@ -378,6 +379,7 @@ class ModelRegistry: descriptor: ModelDescriptor, reqs: RoleRequirements, reference_author_family: Optional[str] = None, + quota_remaining_percent: Optional[float] = None, ) -> Tuple[bool, float, str]: """Evaluate whether model satisfies hard requirements and calculate multidimensional score.""" # 1. Hard Filter: Required capabilities @@ -396,6 +398,8 @@ class ModelRegistry: # 4. Hard Filter: Minimum quality tier if descriptor.quality_tier < reqs.min_quality_tier: return False, 0.0, f"Quality tier {descriptor.quality_tier} < required {reqs.min_quality_tier}" + if quota_remaining_percent is not None and quota_remaining_percent <= 0: + return False, 0.0, "Quota bucket exhausted" # ── Weighted Multi-Dimensional Score ── # Normalized quality: 0.2 to 1.0 @@ -416,12 +420,19 @@ class ModelRegistry: if reqs.diversity_priority > 0 and reference_author_family: div_score = 1.0 if descriptor.family != reference_author_family else 0.2 + # Unknown quota remains neutral; it is never treated as 100% available. + quota_score = 0.5 if quota_remaining_percent is None else max( + 0.0, + min(1.0, quota_remaining_percent / 100.0), + ) + total_score = ( qual_score * reqs.quality_priority + reas_score * reqs.reasoning_priority + lat_score * reqs.latency_priority + cost_score * reqs.cost_priority + div_score * reqs.diversity_priority + + quota_score * reqs.quota_priority ) return True, round(total_score, 4), "Satisfies all capability and quality requirements" diff --git a/src/antigravity_provider/router/profile_manager.py b/src/antigravity_provider/router/profile_manager.py index 306b15f..f1c8fcc 100644 --- a/src/antigravity_provider/router/profile_manager.py +++ b/src/antigravity_provider/router/profile_manager.py @@ -196,13 +196,42 @@ class ProfileAuthManager: @classmethod def save_profile_auth(cls, provider: str, profile_id: str, auth_data: dict) -> Path: - """Save credentials to profile-specific auth.json.""" + """Atomically save credentials and emit a secret-free targeted event.""" pdir = get_profile_dir(profile_id, provider) pdir.mkdir(parents=True, exist_ok=True) auth_file = pdir / "auth.json" - auth_file.write_text(json.dumps(auth_data, indent=2), encoding="utf-8") + existed = auth_file.is_file() + temp_file = pdir / f"auth.json.tmp-{threading.get_ident()}" + temp_file.write_text(json.dumps(auth_data, indent=2), encoding="utf-8") + os.replace(temp_file, auth_file) + + from antigravity_provider.router.event_bus import ( + EVENT_ACCOUNT_ADDED, + EVENT_ACCOUNT_AUTH_CHANGED, + EventBus, + ) + + EventBus.get().publish( + EVENT_ACCOUNT_AUTH_CHANGED if existed else EVENT_ACCOUNT_ADDED, + {"provider": provider, "profile_id": profile_id}, + ) return auth_file + @classmethod + def delete_profile_auth(cls, provider: str, profile_id: str) -> bool: + """Delete one credential file and emit a targeted removal event.""" + auth_file = get_profile_auth_path(provider, profile_id) + if not auth_file.is_file(): + return False + auth_file.unlink() + from antigravity_provider.router.event_bus import EVENT_ACCOUNT_REMOVED, EventBus + + EventBus.get().publish( + EVENT_ACCOUNT_REMOVED, + {"provider": provider, "profile_id": profile_id}, + ) + return True + @classmethod def load_profile_auth(cls, provider: str, profile_id: str) -> Optional[dict]: """Load credentials from profile-specific auth.json or env/auth.json fallback.""" diff --git a/src/antigravity_provider/router/quota_collector.py b/src/antigravity_provider/router/quota_collector.py index 4fa7b74..e7ea635 100644 --- a/src/antigravity_provider/router/quota_collector.py +++ b/src/antigravity_provider/router/quota_collector.py @@ -75,6 +75,26 @@ class AccountQuotaService: return ident return self._resolve_identity(provider, profile_id) + def remaining_for_model(self, provider: str, profile_id: str, model_family: str) -> Optional[float]: + """Return remaining percent for one model pool, or None when the backend has no value.""" + snapshot = self.get_snapshot(provider, profile_id) + if not snapshot or not snapshot.buckets: + return None + family = (model_family or "").lower() + candidates = [ + bucket + for bucket in snapshot.buckets + if not bucket.model_family + or bucket.model_family.lower() in family + or family in bucket.model_family.lower() + ] + measured = [bucket.remaining_percent for bucket in candidates if bucket.remaining_percent is not None] + if measured: + return max(measured) + if any(bucket.status == "exhausted" for bucket in candidates): + return 0.0 + return None + def refresh_account_async(self, provider: str, profile_id: str, on_complete: Optional[Callable[[QuotaSnapshot], None]] = None) -> None: """Fetch fresh quota in a background thread to prevent UI locking.""" def _worker(): @@ -220,10 +240,15 @@ class AccountQuotaService: ) snap.buckets = updated_buckets + snap.source = "runtime_error" with self._cache_lock: self._snapshots[key] = snap logger.info("Runtime quota error recorded for %s model=%s (reset in %ds)", key, model, reset_seconds) + # Local import avoids the state_store -> quota_collector import cycle. + from antigravity_provider.router.state_store import HubStateStore + + HubStateStore.get().apply_delta_quota_updated(provider, profile_id, snap) # ───────────────────────────────────────────────────────────── # IDENTITY RESOLUTION @@ -303,7 +328,7 @@ class AccountQuotaService: remaining_percent=None, period="5h", reset_at=claude_reset_5h, - status="healthy", + status="unknown", ) b_claude_weekly = QuotaBucket( id="antigravity.claude.weekly", @@ -313,7 +338,7 @@ class AccountQuotaService: remaining_percent=None, period="7d", reset_at=weekly_reset, - status="healthy", + status="unknown", ) b_gemini_5h = QuotaBucket( id="antigravity.gemini.5h", @@ -323,7 +348,7 @@ class AccountQuotaService: remaining_percent=None, period="5h", reset_at=gemini_reset_5h, - status="healthy", + status="unknown", ) b_gemini_weekly = QuotaBucket( id="antigravity.gemini.weekly", @@ -333,7 +358,7 @@ class AccountQuotaService: remaining_percent=None, period="7d", reset_at=weekly_reset, - status="healthy", + status="unknown", ) return QuotaSnapshot( @@ -355,7 +380,7 @@ class AccountQuotaService: remaining_percent=None, period="5h", reset_at=now + timedelta(hours=5), - status="healthy", + status="unknown", ) b_weekly = QuotaBucket( id="codex.weekly", @@ -365,7 +390,7 @@ class AccountQuotaService: remaining_percent=None, period="7d", reset_at=now + timedelta(days=7), - status="healthy", + status="unknown", ) return QuotaSnapshot( @@ -386,7 +411,7 @@ class AccountQuotaService: used_percent=None, remaining_percent=None, period="sliding", - status="healthy", + status="unknown", ) b_weekly = QuotaBucket( id="opencode.weekly", @@ -396,7 +421,7 @@ class AccountQuotaService: remaining_percent=None, period="7d", reset_at=now + timedelta(days=7), - status="healthy", + status="unknown", ) b_monthly = QuotaBucket( id="opencode.monthly", @@ -406,7 +431,7 @@ class AccountQuotaService: remaining_percent=None, period="30d", reset_at=now + timedelta(days=30), - status="healthy", + status="unknown", ) return QuotaSnapshot( @@ -428,7 +453,7 @@ class AccountQuotaService: remaining_percent=None, period="5h", reset_at=now + timedelta(hours=5), - status="healthy", + status="unknown", ) b_weekly = QuotaBucket( id="claude.weekly", @@ -438,7 +463,7 @@ class AccountQuotaService: remaining_percent=None, period="7d", reset_at=now + timedelta(days=7), - status="healthy", + status="unknown", ) return QuotaSnapshot( @@ -459,7 +484,7 @@ class AccountQuotaService: used_percent=None, remaining_percent=None, period="7d", - status="healthy", + status="unknown", ) b_chat = QuotaBucket( id="grok.chat", @@ -467,7 +492,7 @@ class AccountQuotaService: model_family="grok", used_percent=None, remaining_percent=None, - status="healthy", + status="unknown", ) b_build = QuotaBucket( id="grok.build", @@ -475,7 +500,7 @@ class AccountQuotaService: model_family="grok", used_percent=None, remaining_percent=None, - status="healthy", + status="unknown", ) b_frequent = QuotaBucket( id="grok.frequent_tasks", @@ -484,7 +509,7 @@ class AccountQuotaService: used_absolute=None, remaining_absolute=None, limit_absolute=10, - status="healthy", + status="unknown", ) b_normal = QuotaBucket( id="grok.normal_tasks", @@ -493,7 +518,7 @@ class AccountQuotaService: used_absolute=None, remaining_absolute=None, limit_absolute=30, - status="healthy", + status="unknown", ) return QuotaSnapshot( @@ -505,19 +530,40 @@ class AccountQuotaService: ) def _generate_baseline_snapshot(self, provider: str, profile_id: str) -> QuotaSnapshot: - """Baseline snapshot when offline or unconfigured.""" + """Truthful offline baseline with provider-specific independent limit pools.""" now = _utc_now() - b = QuotaBucket( - id=f"{provider}.default", - display_name="Основная квота", - used_percent=None, - remaining_percent=None, - status="healthy", + bucket_specs = { + "antigravity": [ + ("antigravity.claude.5h", "Claude 5h", "claude", "5h"), + ("antigravity.gemini.5h", "Gemini 5h", "gemini", "5h"), + ], + "openai-codex": [("codex.primary.weekly", "Codex Weekly", "gpt", "7d")], + "codex": [("codex.primary.weekly", "Codex Weekly", "gpt", "7d")], + "claude": [("claude.session.5h", "Claude 5h", "claude", "5h")], + "anthropic": [("claude.session.5h", "Claude 5h", "claude", "5h")], + "grok": [("grok.frequent_tasks", "Grok 2h", "grok", "2h")], + "xai": [("grok.frequent_tasks", "Grok 2h", "grok", "2h")], + "opencode-go": [("opencode.tasks", "OpenCode Tasks", "opencode", "30d")], + "opencode": [("opencode.tasks", "OpenCode Tasks", "opencode", "30d")], + } + specs = bucket_specs.get( + provider, + [(f"{provider}.default", "Основная квота", None, None)], ) + buckets = [ + QuotaBucket( + id=bucket_id, + display_name=display_name, + model_family=model_family, + period=period, + status="unknown", + ) + for bucket_id, display_name, model_family, period in specs + ] return QuotaSnapshot( account_id=profile_id, provider=provider, - buckets=[b], + buckets=buckets, fetched_at=now, source="baseline", ) diff --git a/src/antigravity_provider/router/router_engine.py b/src/antigravity_provider/router/router_engine.py index 37ebc91..9f74eaa 100644 --- a/src/antigravity_provider/router/router_engine.py +++ b/src/antigravity_provider/router/router_engine.py @@ -149,7 +149,18 @@ class RouterEngine: for m_candidate in viable_models: m_desc = registry.get_model(m_candidate) if m_desc: - ok, score, _ = registry.evaluate_model_score(m_desc, role_reqs) + from .quota_collector import AccountQuotaService + + quota_remaining = AccountQuotaService.get().remaining_for_model( + pconfig.provider, + pid, + m_desc.family, + ) + ok, score, _ = registry.evaluate_model_score( + m_desc, + role_reqs, + quota_remaining_percent=quota_remaining, + ) if ok: scored_candidates.append((score, m_candidate)) else: diff --git a/src/antigravity_provider/router/scheduler.py b/src/antigravity_provider/router/scheduler.py index 26a3b8b..7fb5105 100644 --- a/src/antigravity_provider/router/scheduler.py +++ b/src/antigravity_provider/router/scheduler.py @@ -19,6 +19,8 @@ from typing import Any, Callable, Dict, List, Optional, Set from antigravity_provider.router.event_bus import ( EventBus, + EVENT_ACCOUNT_ADDED, + EVENT_ACCOUNT_AUTH_CHANGED, EVENT_REFRESH_STARTED, EVENT_REFRESH_COMPLETED, EVENT_REFRESH_FAILED, @@ -85,6 +87,17 @@ class HermesRefreshScheduler: self.tasks_deduplicated_total: int = 0 self._init_default_tasks() + EventBus.get().subscribe(EVENT_ACCOUNT_ADDED, self._on_account_event) + EventBus.get().subscribe(EVENT_ACCOUNT_AUTH_CHANGED, self._on_account_event) + + def _on_account_event(self, _event_name: str, payload: Any) -> None: + """Refresh only the account mentioned by an OAuth/auth lifecycle event.""" + if not isinstance(payload, dict): + return + provider = payload.get("provider") + profile_id = payload.get("profile_id") + if provider and profile_id: + self.trigger_refresh_account(str(provider), str(profile_id)) @classmethod def get(cls) -> HermesRefreshScheduler: @@ -235,14 +248,32 @@ class HermesRefreshScheduler: if task.scope == "single" and task.profile_id: status = uh_service.get_profile_status(task.provider, task.profile_id) if status.get("authenticated"): - quota_service.refresh_account_async(task.provider, task.profile_id) + quota_snapshot = quota_service.fetch_account_quota( + task.provider, + task.profile_id, + force=True, + ) + HubStateStore.get().apply_delta_quota_updated( + task.provider, + task.profile_id, + quota_snapshot, + ) elif task.scope in ("full", "current"): # Refresh quota snapshots for configured accounts of this provider profs = uh_service.get_cached_profiles().get(task.provider, []) for p in profs: if p.auth_state == "AUTHENTICATED": - quota_service.refresh_account_async(task.provider, p.profile_id) + quota_snapshot = quota_service.fetch_account_quota( + task.provider, + p.profile_id, + force=True, + ) + HubStateStore.get().apply_delta_quota_updated( + task.provider, + p.profile_id, + quota_snapshot, + ) # Rebuild unified snapshot HubStateStore.get().refresh(force_scan=True, seq=seq) @@ -279,7 +310,8 @@ class HermesRefreshScheduler: def _worker(): try: quota_service = AccountQuotaService.get() - quota_service.refresh_account_async(provider, profile_id) + quota_snapshot = quota_service.fetch_account_quota(provider, profile_id, force=True) + HubStateStore.get().apply_delta_quota_updated(provider, profile_id, quota_snapshot) HubStateStore.get().apply_delta_account_updated(profile_id) finally: with self._lock: @@ -304,8 +336,13 @@ class HermesRefreshScheduler: def _worker(): try: - AccountQuotaService.get().refresh_all_accounts_async() - HubStateStore.get().refresh(force_scan=True) + quota_service = AccountQuotaService.get() + results = quota_service.fetch_all_configured(force=True) + store = HubStateStore.get() + for key, quota_snapshot in results.items(): + provider, profile_id = key.split(":", 1) + store.apply_delta_quota_updated(provider, profile_id, quota_snapshot) + store.refresh(force_scan=True) finally: with self._lock: self._in_flight_refreshes.pop(key, None) diff --git a/src/antigravity_provider/router/skills/skill_registry.py b/src/antigravity_provider/router/skills/skill_registry.py deleted file mode 100644 index 7b561ca..0000000 --- a/src/antigravity_provider/router/skills/skill_registry.py +++ /dev/null @@ -1,130 +0,0 @@ -"""Unified Skill Registry supporting Hermes Agent, Google Antigravity, OpenAI Codex, and OpenCode.""" -from __future__ import annotations - -import json -from dataclasses import asdict, dataclass, field -from pathlib import Path -from typing import Any, Dict, List, Optional, Tuple - - -@dataclass -class SkillParameter: - name: str - type_name: str - description: str - required: bool = True - default: Optional[Any] = None - - -@dataclass -class UnifiedSkill: - skill_id: str - name: str - description: str - parameters: List[SkillParameter] = field(default_factory=list) - tags: List[str] = field(default_factory=list) - supported_providers: List[str] = field(default_factory=lambda: ["antigravity", "openai-codex", "opencode-go"]) - requires_approval: bool = False - - -class UnifiedSkillRegistry: - """Central registry normalizing skills across multi-provider formats.""" - - _instance: Optional[UnifiedSkillRegistry] = None - - def __init__(self): - self._skills: Dict[str, UnifiedSkill] = {} - self._register_default_skills() - - @classmethod - def get(cls) -> UnifiedSkillRegistry: - if cls._instance is None: - cls._instance = cls() - return cls._instance - - def register(self, skill: UnifiedSkill): - self._skills[skill.skill_id] = skill - - def get_skill(self, skill_id: str) -> Optional[UnifiedSkill]: - return self._skills.get(skill_id) - - def list_skills(self, provider: Optional[str] = None) -> List[UnifiedSkill]: - if not provider: - return list(self._skills.values()) - return [s for s in self._skills.values() if provider in s.supported_providers] - - def to_provider_schema(self, skill_id: str, provider: str) -> Dict[str, Any]: - """Translate unified skill into provider-specific tool call declaration.""" - skill = self._skills.get(skill_id) - if not skill: - raise KeyError(f"Skill '{skill_id}' not found") - - props = {} - required = [] - for p in skill.parameters: - props[p.name] = { - "type": p.type_name, - "description": p.description, - } - if p.required: - required.append(p.name) - - if "antigravity" in provider or "google" in provider: - return { - "name": skill.name, - "description": skill.description, - "parameters": { - "type": "OBJECT", - "properties": props, - "required": required, - }, - } - else: # OpenAI / OpenCode format - return { - "type": "function", - "function": { - "name": skill.name, - "description": skill.description, - "parameters": { - "type": "object", - "properties": props, - "required": required, - }, - }, - } - - def _register_default_skills(self): - self.register( - UnifiedSkill( - skill_id="search_web", - name="search_web", - description="Performs internet web search queries", - parameters=[ - SkillParameter(name="query", type_name="string", description="The search term"), - ], - tags=["search", "web", "research"], - ) - ) - self.register( - UnifiedSkill( - skill_id="run_command", - name="run_command", - description="Execute commands safely in the target terminal shell", - parameters=[ - SkillParameter(name="CommandLine", type_name="string", description="Exact command line string"), - SkillParameter(name="Cwd", type_name="string", description="Working directory"), - ], - tags=["terminal", "execution"], - ) - ) - self.register( - UnifiedSkill( - skill_id="view_file", - name="view_file", - description="View contents of a file from filesystem", - parameters=[ - SkillParameter(name="AbsolutePath", type_name="string", description="Target file path"), - ], - tags=["file", "read"], - ) - ) diff --git a/src/antigravity_provider/router/state_store.py b/src/antigravity_provider/router/state_store.py index 0558b71..c2a018a 100644 --- a/src/antigravity_provider/router/state_store.py +++ b/src/antigravity_provider/router/state_store.py @@ -8,12 +8,15 @@ from __future__ import annotations import logging import threading import time -from dataclasses import dataclass, field +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, @@ -40,6 +43,7 @@ logger = logging.getLogger("hermes.router.state_store") 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] @@ -104,53 +108,58 @@ class HubStateStore: return self.refresh(force_scan=False) def refresh(self, force_scan: bool = True, seq: Optional[int] = None) -> HubSnapshot: - """Execute a single unified state build cycle and publish an updated 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: - if seq is not None and seq < self._latest_applied_seq: - logger.warning("Rejecting stale refresh result (seq %d < applied %d)", seq, self._latest_applied_seq) + 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.refresh_runs_total += 1 - if seq is not None: - self._latest_applied_seq = max(self._latest_applied_seq, seq) - - t0 = time.time() + self._latest_applied_seq = request_seq self._generation += 1 gen = self._generation - - # Single unified scan - uh_service = UnifiedHealthService.get() - profiles_by_prov = uh_service.scan_all(force=force_scan) - - all_profs: Dict[str, ProfileViewModel] = {} - for prov, profs in profiles_by_prov.items(): - for p in profs: - all_profs[p.profile_id] = p - - 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() - - # Quotas map - quota_service = AccountQuotaService.get() - quotas_map: Dict[str, Any] = {} - for pid, p in all_profs.items(): - if p.auth_state == "AUTHENTICATED": - quotas_map[pid] = quota_service.get_snapshot(p.provider, pid) - metrics = { "generation": gen, + "seq": request_seq, "duration_ms": round((time.time() - t0) * 1000, 2), "total_profiles": len(all_profs), - "authenticated_profiles": sum(1 for p in all_profs.values() if p.auth_state == "AUTHENTICATED"), + "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, @@ -162,17 +171,20 @@ class HubStateStore: 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, "duration_ms": metrics["duration_ms"]}) + 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={}, @@ -185,32 +197,154 @@ class HubStateStore: is_stale=True, ) - def apply_delta_account_updated(self, profile_id: str) -> None: - """Apply targeted account delta update and notify UI without global scan.""" + def _apply_profile_delta(self, profile: ProfileViewModel) -> HubSnapshot: + """Copy-on-write replacement of exactly one profile in the snapshot.""" with self._lock: - self.account_updates_total += 1 - # Invalidate cached view model for targeted profile - uh_service = UnifiedHealthService.get() - with uh_service._lock: - uh_service._cached_profiles.pop(profile_id, None) + 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 - # Refresh snapshot and notify - snap = self.refresh(force_scan=False) - updated_profile = snap.get_profile(profile_id) - if updated_profile: - EventBus.get().publish(EVENT_ACCOUNT_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": updated_profile, - "generation": snap.generation, - }) + "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, - }) + EventBus.get().publish( + EVENT_QUOTA_UPDATED, + { + "provider": provider, + "profile_id": profile_id, + "quota_snapshot": quota_snap, + "generation": updated.generation, + "seq": updated.seq, + }, + ) diff --git a/src/antigravity_provider/router/supervisor/lifecycle_supervisor.py b/src/antigravity_provider/router/supervisor/lifecycle_supervisor.py deleted file mode 100644 index 9ceaaa5..0000000 --- a/src/antigravity_provider/router/supervisor/lifecycle_supervisor.py +++ /dev/null @@ -1,221 +0,0 @@ -"""Lifecycle Supervisor & Process Registry with strict process ownership metadata. - -Invariant: Never use generic killall / taskkill python. Process cleanup only targets -verified child processes registered with an active UUID lease. -""" -from __future__ import annotations - -import json -import logging -import os -import psutil -import threading -import time -import uuid -from dataclasses import asdict, dataclass, field -from pathlib import Path -from typing import Any, Dict, List, Optional, Set - -from antigravity_provider.router.unified_health import EventLogService - -logger = logging.getLogger("hermes.router.supervisor") - - -@dataclass -class ProcessEntry: - process_uuid: str - pid: int - name: str - cmdline: List[str] - owner_app: str - spawn_time: float - lease_ttl_sec: float - last_heartbeat: float - status: str = "running" # running | stopped | dead | zombie - - -@dataclass -class LeaseRecord: - lease_id: str - profile_id: str - consumer: str - acquired_at: float - ttl_sec: float - expires_at: float - metadata: Dict[str, Any] = field(default_factory=dict) - - -class LifecycleSupervisor: - """Central supervisor managing process lifecycle, heartbeat, and ownership leases.""" - - _instance: Optional[LifecycleSupervisor] = None - _lock = threading.RLock() - - def __init__(self, state_dir: Optional[Path] = None): - if state_dir is None: - local_app = os.environ.get("LOCALAPPDATA", "") - state_dir = Path(local_app) / "hermes" / "supervisor" - self.state_dir = state_dir - self.state_dir.mkdir(parents=True, exist_ok=True) - self.registry_file = self.state_dir / "process_registry.json" - - self._processes: Dict[str, ProcessEntry] = {} - self._leases: Dict[str, LeaseRecord] = {} - self._running = True - self._supervisor_thread: Optional[threading.Thread] = None - - self._load_registry() - self._start_supervisor() - - @classmethod - def get(cls) -> LifecycleSupervisor: - if cls._instance is None: - cls._instance = cls() - return cls._instance - - def register_process( - self, - pid: int, - name: str, - cmdline: List[str], - owner_app: str = "HermesHub", - ttl_sec: float = 300.0, - ) -> ProcessEntry: - """Register a new child process with explicit ownership metadata.""" - with self._lock: - p_uuid = str(uuid.uuid4()) - now = time.time() - entry = ProcessEntry( - process_uuid=p_uuid, - pid=pid, - name=name, - cmdline=cmdline, - owner_app=owner_app, - spawn_time=now, - lease_ttl_sec=ttl_sec, - last_heartbeat=now, - status="running", - ) - self._processes[p_uuid] = entry - self._save_registry() - EventLogService.get().log("system", f"Процесс зарегистрирован: {name} (PID: {pid}, UUID: {p_uuid[:8]})", level="info") - return entry - - def heartbeat(self, process_uuid: str) -> bool: - """Update heartbeat for an active process.""" - with self._lock: - if process_uuid in self._processes: - self._processes[process_uuid].last_heartbeat = time.time() - return True - return False - - def acquire_lease(self, profile_id: str, consumer: str, ttl_sec: float = 60.0) -> Optional[LeaseRecord]: - """Acquire a temporary exclusive usage lease for a profile.""" - with self._lock: - now = time.time() - # Clean expired leases - self._leases = {lid: l for lid, l in self._leases.items() if l.expires_at > now} - - # Check existing active lease - for l in self._leases.values(): - if l.profile_id == profile_id and l.expires_at > now: - return None # Profile currently leased - - lease_id = str(uuid.uuid4()) - record = LeaseRecord( - lease_id=lease_id, - profile_id=profile_id, - consumer=consumer, - acquired_at=now, - ttl_sec=ttl_sec, - expires_at=now + ttl_sec, - ) - self._leases[lease_id] = record - return record - - def release_lease(self, lease_id: str) -> bool: - with self._lock: - if lease_id in self._leases: - del self._leases[lease_id] - return True - return False - - def terminate_owned_process(self, process_uuid: str) -> bool: - """Safely terminate a verified owned process without collateral damage.""" - with self._lock: - entry = self._processes.get(process_uuid) - if not entry: - return False - - pid = entry.pid - try: - if psutil.pid_exists(pid): - p = psutil.Process(pid) - # Verify command line matches ownership record - p_cmd = p.cmdline() if hasattr(p, "cmdline") else [] - if entry.name.lower() in p.name().lower() or (p_cmd and p_cmd[0] in entry.cmdline[0]): - p.terminate() - try: - p.wait(timeout=3) - except psutil.TimeoutExpired: - p.kill() - except (psutil.NoSuchProcess, psutil.AccessDenied): - pass - - entry.status = "stopped" - self._save_registry() - EventLogService.get().log("system", f"Процесс остановлен: {entry.name} (PID: {pid})", level="info") - return True - - def cleanup_expired_processes(self) -> int: - """Clean up strictly owned orphaned processes that missed heartbeats past TTL.""" - with self._lock: - now = time.time() - cleaned = 0 - for p_uuid, entry in list(self._processes.items()): - if entry.status == "running": - if now - entry.last_heartbeat > entry.lease_ttl_sec: - # Process timed out - self.terminate_owned_process(p_uuid) - cleaned += 1 - return cleaned - - def shutdown_all_owned(self): - """Clean shutdown of all registered child processes.""" - with self._lock: - self._running = False - for p_uuid, entry in list(self._processes.items()): - if entry.status == "running": - self.terminate_owned_process(p_uuid) - - def _start_supervisor(self): - def _loop(): - while self._running: - time.sleep(5) - try: - self.cleanup_expired_processes() - except Exception: - pass - - self._supervisor_thread = threading.Thread(target=_loop, daemon=True) - self._supervisor_thread.start() - - def _save_registry(self): - try: - data = {k: asdict(v) for k, v in self._processes.items()} - self.registry_file.write_text(json.dumps(data, indent=2), encoding="utf-8") - except Exception: - pass - - def _load_registry(self): - if self.registry_file.exists(): - try: - data = json.loads(self.registry_file.read_text(encoding="utf-8")) - for k, v in data.items(): - # Validate if process still running - p_entry = ProcessEntry(**v) - if p_entry.status == "running" and not psutil.pid_exists(p_entry.pid): - p_entry.status = "dead" - self._processes[k] = p_entry - except Exception: - pass diff --git a/tests/test_data_truthfulness_and_oauth_security.py b/tests/test_data_truthfulness_and_oauth_security.py index 195dc7f..8b6d6dc 100644 --- a/tests/test_data_truthfulness_and_oauth_security.py +++ b/tests/test_data_truthfulness_and_oauth_security.py @@ -47,16 +47,16 @@ def test_quota_collector_never_fakes_api_source_without_network(): def test_quota_bucket_formatted_remaining_honesty(): - """P0-1: Bucket formatted_remaining returns honest availability when percentages are None.""" + """P0-1: Missing values are unknown, not evidence that quota is available.""" b = QuotaBucket( id="test.bucket", display_name="Test Bucket", used_percent=None, remaining_percent=None, - status="healthy", + status="unknown", ) - assert b.formatted_remaining() == "Доступна" - assert b.status == "healthy" + assert b.formatted_remaining() == "Н/Д" + assert b.status == "unknown" # ── TEST P0-2: OAuth Fail-Closed & DEV_MODE Gating ── diff --git a/tests/test_import_source_contract.py b/tests/test_import_source_contract.py new file mode 100644 index 0000000..8f86b5a --- /dev/null +++ b/tests/test_import_source_contract.py @@ -0,0 +1,18 @@ +"""The repository and installed plugin must never merge into one namespace.""" +from __future__ import annotations + +from pathlib import Path + +import antigravity_provider +import antigravity_provider.runtime as runtime + + +def test_provider_is_regular_package_from_repository() -> None: + repo_src = (Path(__file__).resolve().parent.parent / "src").resolve() + package_file = Path(antigravity_provider.__file__).resolve() + runtime_file = Path(runtime.__file__).resolve() + + assert package_file.is_relative_to(repo_src) + assert runtime_file.is_relative_to(repo_src) + assert hasattr(runtime, "format_antigravity_error") + assert list(antigravity_provider.__path__) == [str(package_file.parent)] diff --git a/tests/test_roadmap_features.py b/tests/test_roadmap_features.py index 6f3d67a..1936428 100644 --- a/tests/test_roadmap_features.py +++ b/tests/test_roadmap_features.py @@ -1,12 +1,6 @@ -"""Tests for Roadmap Features (Lifecycle Supervisor, Policies, Skill Registry, Capability Matrix).""" -import pytest -import time -from pathlib import Path +"""Tests for active roadmap policies and normalized health.""" -from antigravity_provider.router.supervisor.lifecycle_supervisor import LifecycleSupervisor from antigravity_provider.router.supervisor.policies import PolicyEnforcer, WebPolicyConfig, ToolPolicyConfig -from antigravity_provider.router.skills.skill_registry import UnifiedSkillRegistry, UnifiedSkill, SkillParameter -from antigravity_provider.router.capability.capability_matrix import CapabilityMatrix from antigravity_provider.router.unified_health import ( UnifiedHealthService, STATUS_HEALTHY, @@ -29,37 +23,6 @@ def test_status_resolver_unconfigured_account(): assert p.health_state != "quota_exhausted" -def test_lifecycle_supervisor_process_registration_and_lease(tmp_path): - supervisor = LifecycleSupervisor(state_dir=tmp_path) - entry = supervisor.register_process( - pid=12345, - name="test_worker", - cmdline=["python", "-m", "worker"], - owner_app="HermesHub", - ttl_sec=10.0, - ) - assert entry.pid == 12345 - assert entry.status == "running" - - # Heartbeat - assert supervisor.heartbeat(entry.process_uuid) is True - - # Acquire lease - lease1 = supervisor.acquire_lease("antigravity-1", "orchestrator", ttl_sec=5.0) - assert lease1 is not None - assert lease1.profile_id == "antigravity-1" - - # Concurrent lease for same profile should fail - lease2 = supervisor.acquire_lease("antigravity-1", "coder", ttl_sec=5.0) - assert lease2 is None - - # Release lease - assert supervisor.release_lease(lease1.lease_id) is True - # Now coder can acquire - lease3 = supervisor.acquire_lease("antigravity-1", "coder", ttl_sec=5.0) - assert lease3 is not None - - def test_web_policy_enforcement(): enforcer = PolicyEnforcer() @@ -89,28 +52,3 @@ def test_tool_policy_enforcement(): ok, _ = enforcer.validate_tool_execution("run_command", {"CommandLine": "format C: /y"}) assert ok is False - -def test_unified_skill_registry(): - reg = UnifiedSkillRegistry.get() - skills = reg.list_skills() - assert len(skills) >= 3 - - # Check provider translation - ag_schema = reg.to_provider_schema("search_web", "antigravity") - assert ag_schema["name"] == "search_web" - assert "parameters" in ag_schema - assert ag_schema["parameters"]["type"] == "OBJECT" - - openai_schema = reg.to_provider_schema("search_web", "openai-codex") - assert openai_schema["type"] == "function" - assert openai_schema["function"]["name"] == "search_web" - - -def test_capability_matrix_role_selection(): - matrix = CapabilityMatrix.get() - # Find models supporting reasoning and tools - candidates = matrix.find_best_model_for_role(required_tools=True, required_reasoning=True) - assert len(candidates) >= 2 - model_ids = [c.model_id for c in candidates] - assert "gemini-2.5-pro" in model_ids - assert "gpt-5.3-codex" in model_ids diff --git a/tests/test_state_layer_contract.py b/tests/test_state_layer_contract.py new file mode 100644 index 0000000..711dc09 --- /dev/null +++ b/tests/test_state_layer_contract.py @@ -0,0 +1,213 @@ +"""Task A regressions for sequence ordering and point updates.""" +from __future__ import annotations + +import threading +import time +from unittest.mock import MagicMock + +import pytest + +from antigravity_provider.router.event_bus import ( + EventBus, + EVENT_ACCOUNT_ADDED, + EVENT_ACCOUNT_REMOVED, + EVENT_QUOTA_UPDATED, +) +from antigravity_provider.router.profile_manager import ProfileAuthManager +from antigravity_provider.router.scheduler import HermesRefreshScheduler +from antigravity_provider.router.state_store import HubSnapshot, HubStateStore +from antigravity_provider.router.model_registry import ModelRegistry +from antigravity_provider.router.unified_health import ProfileViewModel, SystemReadiness + + +def _readiness() -> SystemReadiness: + return SystemReadiness( + state="limited", + title_ru="Тест", + summary_ru="", + roles_ready_count=0, + total_roles=0, + accounts_connected_count=0, + total_accounts=0, + providers_ready_count=0, + total_providers=0, + ) + + +def _profile(profile_id: str) -> ProfileViewModel: + return ProfileViewModel( + profile_id=profile_id, + display_name=profile_id, + account_identity=f"{profile_id}@example.test", + provider="antigravity", + provider_display_name="Antigravity", + assigned_roles=[], + primary_role=None, + is_main_account=False, + is_main_orchestrator=False, + auth_state="AUTHENTICATED", + health_state="healthy", + health_label_ru="Работает", + model_states={}, + cooldown_remaining_sec=0, + last_checked_at=None, + enabled=True, + is_cold_spare=False, + is_empty_slot=False, + ) + + +@pytest.mark.unit +def test_late_refresh_result_cannot_overwrite_newer_seq(monkeypatch: pytest.MonkeyPatch) -> None: + slow_started = threading.Event() + release_slow = threading.Event() + + class FakeHealth: + def scan_all(self, force: bool = False): + if threading.current_thread().name == "slow-refresh": + slow_started.set() + assert release_slow.wait(2) + return {} + + def get_system_readiness(self): + return _readiness() + + def get_agent_view_models(self): + return [] + + def get_provider_summaries(self): + return [] + + def get_routing_pipelines(self): + return {} + + monkeypatch.setattr( + "antigravity_provider.router.state_store.UnifiedHealthService.get", + lambda: FakeHealth(), + ) + + store = HubStateStore() + slow = threading.Thread(target=lambda: store.refresh(seq=1), name="slow-refresh") + slow.start() + assert slow_started.wait(1) + fast_snapshot = store.refresh(seq=2) + release_slow.set() + slow.join(timeout=2) + + assert fast_snapshot.seq == 2 + assert store.get_snapshot().seq == 2 + assert store.get_snapshot().generation == 1 + assert store.refresh_skipped_total == 1 + + +@pytest.mark.unit +def test_quota_delta_updates_only_target_profile_and_emits_key() -> None: + profile_a = _profile("account-a") + profile_b = _profile("account-b") + store = HubStateStore() + store._current_snapshot = HubSnapshot( + generation=1, + seq=1, + timestamp=time.time(), + profiles_by_provider={"antigravity": [profile_a, profile_b]}, + all_profiles={"account-a": profile_a, "account-b": profile_b}, + readiness=_readiness(), + agents=[], + providers=[], + routing={}, + quotas={}, + ) + received = [] + bus = EventBus.get() + bus.subscribe(EVENT_QUOTA_UPDATED, lambda _name, payload: received.append(payload)) + try: + quota = object() + store.apply_delta_quota_updated("antigravity", "account-a", quota) + finally: + bus._listeners.clear() + + snapshot = store.get_snapshot() + assert snapshot.all_profiles["account-b"] is profile_b + assert snapshot.all_profiles["account-a"] is not profile_a + assert snapshot.quotas["account-a"] is quota + assert received[-1]["profile_id"] == "account-a" + assert received[-1]["seq"] == snapshot.seq + + +@pytest.mark.unit +def test_single_scheduler_refresh_waits_for_quota_before_delta(monkeypatch: pytest.MonkeyPatch) -> None: + order: list[str] = [] + quota_service = MagicMock() + quota_service.fetch_account_quota.side_effect = lambda *_args, **_kwargs: order.append("fetch") or "quota" + store = MagicMock() + store.apply_delta_quota_updated.side_effect = lambda *_args: order.append("quota_delta") + store.apply_delta_account_updated.side_effect = lambda *_args: order.append("account_delta") + monkeypatch.setattr( + "antigravity_provider.router.scheduler.AccountQuotaService.get", + lambda: quota_service, + ) + monkeypatch.setattr("antigravity_provider.router.scheduler.HubStateStore.get", lambda: store) + + done = threading.Event() + scheduler = HermesRefreshScheduler() + scheduler.trigger_refresh_account("antigravity", "account-a", on_complete=done.set) + assert done.wait(2) + assert order == ["fetch", "quota_delta", "account_delta"] + + +@pytest.mark.unit +def test_model_score_rejects_exhausted_pool_and_prefers_more_quota() -> None: + registry = ModelRegistry.get() + descriptor = registry.get_model("gemini-2.5-pro") + requirements = registry.get_role_requirements("research") + assert descriptor is not None + + ok_empty, _, reason = registry.evaluate_model_score( + descriptor, + requirements, + quota_remaining_percent=0, + ) + ok_low, score_low, _ = registry.evaluate_model_score( + descriptor, + requirements, + quota_remaining_percent=10, + ) + ok_high, score_high, _ = registry.evaluate_model_score( + descriptor, + requirements, + quota_remaining_percent=90, + ) + + assert not ok_empty and reason == "Quota bucket exhausted" + assert ok_low and ok_high + assert score_high > score_low + + +@pytest.mark.unit +def test_auth_storage_emits_secret_free_targeted_lifecycle_events( + tmp_path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("HERMES_HOME", str(tmp_path / "hermes")) + bus = EventBus.get() + bus._listeners.clear() + received: list[tuple[str, dict]] = [] + bus.subscribe("*", lambda name, payload: received.append((name, payload))) + + ProfileAuthManager.save_profile_auth( + "antigravity", + "account-a", + {"access_token": "must-not-leak", "email": "person@example.test"}, + ) + assert received[-1] == ( + EVENT_ACCOUNT_ADDED, + {"provider": "antigravity", "profile_id": "account-a"}, + ) + assert "must-not-leak" not in repr(received) + + assert ProfileAuthManager.delete_profile_auth("antigravity", "account-a") + assert received[-1] == ( + EVENT_ACCOUNT_REMOVED, + {"provider": "antigravity", "profile_id": "account-a"}, + ) + bus._listeners.clear()