diff --git a/.env.example b/.env.example index 01c389a..23132f8 100644 --- a/.env.example +++ b/.env.example @@ -27,11 +27,25 @@ JWT_SECRET=change-me ENCRYPTION_KEY=change-me-32-bytes # --- AI models (hidden from customers) ------------------------------------- -# mock | openai_compatible +# mock | openai_compatible | agy AI_DEFAULT_PROVIDER=mock OPENAI_COMPATIBLE_BASE_URL= OPENAI_COMPATIBLE_API_KEY= +# --- agy (Google Antigravity) ---------------------------------------------- +# Only used when AI_DEFAULT_PROVIDER=agy. There are no API keys: accounts are +# Google identities logged in via OAuth from the admin console +# (/admin/accounts → افزودن اکانت agy). Any number of accounts can be added — +# the rotation pool spreads load across them and benches a rate-limited one +# until its quota window resets. AGY_CLIENT_ID/SECRET default to the public +# Antigravity CLI application; override them to use your own Google OAuth +# client. AGY_REDIRECT_URI must stay a loopback URL. +AGY_CLIENT_ID= +AGY_CLIENT_SECRET= +AGY_REDIRECT_URI=http://localhost:51121/oauth-callback +# Seconds an account is benched after a 429 when Google sends no Retry-After. +AGY_429_COOLDOWN_SECONDS=900 + # --- Git provider ---------------------------------------------------------- # mock | gitea GIT_PROVIDER=mock diff --git a/.gitignore b/.gitignore index 1ec7afe..9197675 100644 --- a/.gitignore +++ b/.gitignore @@ -36,6 +36,10 @@ venv/ *.log logs/ +# pending agy OAuth login challenge (contains a PKCE verifier — never commit) +.agy-login-state.json +apps/api/.agy-login-state.json + # editor / os .DS_Store .idea/ diff --git a/README.md b/README.md index 4de69f4..1d5234e 100644 --- a/README.md +++ b/README.md @@ -217,7 +217,7 @@ web/ | قابلیت | env | مقادیر | |---|---|---| -| مدل | `AI_DEFAULT_PROVIDER` | `mock` \| `openai_compatible` | +| مدل | `AI_DEFAULT_PROVIDER` | `mock` \| `openai_compatible` \| `agy` | | Git | `GIT_PROVIDER` | `mock` \| `gitea` | | دامنهٔ تست | `DOMAIN_ROUTER` | `mock` \| `caddy` | | Sandbox | `SANDBOX_PROVIDER` | `mock` \| `daytona` \| `e2b` | @@ -226,6 +226,25 @@ web/ اگر credential واقعی نداریم، **قابلیت حذف نمی‌شود**: مدل دیتابیس، interface، UI و adapter mock کامل ساخته می‌شود و adapter واقعی پشت همین flag آماده می‌ماند. +### ارائه‌دهندهٔ agy (Google Antigravity) + +agy کلید API ندارد؛ **اکانت**‌ها هویت Google هستند و از طریق OAuth وارد +می‌شوند. برای فعال‌سازی: + +1. `AI_DEFAULT_PROVIDER=agy` را در `.env` بگذارید و `pnpm db:seed` را دوباره + اجرا کنید (کتالوگ مدل‌های agy فعال و کاتالوگ mock غیرفعال می‌شود). +2. با حساب مالک سامانه وارد کنسول مدیریت شوید و از + `/admin/accounts` → «افزودن اکانت agy» شروع کنید: لینک احراز هویت گوگل + باز می‌شود، بعد از تأیید، URL بازگشت را از نوار آدرس مرورگر کپی و Paste + کنید (صفحهٔ مقصد عمداً بارگذاری نمی‌شود؛ پارامترها همان‌جا هستند). +3. هر تعداد اکانت Google که بخواهید اضافه کنید. مدل‌ها را در + `/admin/models` در صورت نیاز فعال کنید. + +پول چرخشی، بار را بین اکانت‌ها پخش می‌کند و اکانتی که ۴۲۹ گرفته تا پایان +پنجرهٔ سهمیهٔ آن کنار گذاشته می‌شود (هم در حافظه و هم در دیتابیس، تا workerها +هم همان تصویر را ببینند). refresh tokenها فقط encrypted at rest ذخیره +می‌شوند و هیچ‌وقت از API برنمی‌گردند. + --- ## ۹. Workerهای پس‌زمینه diff --git a/apps/api/alembic/versions/b1c2d3e4f5a6_add_agy_account_columns.py b/apps/api/alembic/versions/b1c2d3e4f5a6_add_agy_account_columns.py new file mode 100644 index 0000000..25422e9 --- /dev/null +++ b/apps/api/alembic/versions/b1c2d3e4f5a6_add_agy_account_columns.py @@ -0,0 +1,57 @@ +"""add agy account columns to ai_model_credentials + +Turns ``ai_model_credentials`` into a store that can hold *accounts* for the +agy (Google Antigravity) provider, not just API keys: ``email`` / ``project_id`` +carry the account identity, ``status`` + ``cooldown_until`` drive the rotation +pool (skip a rate-limited account until its reset time), and +``last_error`` / ``last_used_at`` support the admin console. + +Every new NOT NULL column carries a ``server_default`` so existing rows +backfill without a table rewrite and stay insertable from plain SQL. + +Revision ID: b1c2d3e4f5a6 +Revises: f5ba57af2906 +Create Date: 2026-10-13 02:10:00.000000 +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import op + + +revision: str = 'b1c2d3e4f5a6' +down_revision: str | None = 'f5ba57af2906' +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +#: ``(name, type, server_default)``. ``None`` default means nullable, no default. +_NEW_COLUMNS: tuple[tuple[str, sa.types.TypeEngine, str | None], ...] = ( + ("email", sa.String(length=320), "''"), + ("project_id", sa.String(length=160), "''"), + ("status", sa.String(length=16), "'active'"), + ("cooldown_until", sa.DateTime(timezone=True), None), + ("last_error", sa.String(length=512), "''"), + ("last_used_at", sa.DateTime(timezone=True), None), +) + + +def upgrade() -> None: + for name, column_type, default in _NEW_COLUMNS: + op.add_column( + "ai_model_credentials", + sa.Column( + name, + column_type, + nullable=default is None, + server_default=None if default is None else sa.text(default), + ), + ) + + +def downgrade() -> None: + for name, _type, _default in _NEW_COLUMNS: + op.drop_column("ai_model_credentials", name) diff --git a/apps/api/app/application/ports/model_provider.py b/apps/api/app/application/ports/model_provider.py index 0a19d87..d51523b 100644 --- a/apps/api/app/application/ports/model_provider.py +++ b/apps/api/app/application/ports/model_provider.py @@ -44,13 +44,21 @@ class ModelProviderPort(Protocol): self, *, model: AiModel, - credential: ModelCredential | None, + credential: ModelCredential | None = None, + credentials: list[ModelCredential] | None = None, messages: list[ChatMessage], system: str = "", temperature: float = 0.2, max_tokens: int = 2048, ) -> ModelResponse: - """Run one completion. Must raise ``ProviderUnavailableError`` on failure.""" + """Run one completion. Must raise ``ProviderUnavailableError`` on failure. + + ``credential`` is the single active credential (the only input key-based + providers need). ``credentials`` is the full usable pool: providers whose + accounts rotate (agy) accept it and fan out across the set, keeping one + account's rate limit from becoming everyone's outage; providers without + a pool ignore it. + """ ... async def health(self) -> ModelHealth: diff --git a/apps/api/app/application/services/workforce_service.py b/apps/api/app/application/services/workforce_service.py index 0175b0d..c6441f9 100644 --- a/apps/api/app/application/services/workforce_service.py +++ b/apps/api/app/application/services/workforce_service.py @@ -714,10 +714,17 @@ class WorkforceService: ) -> str: """Run one completion and record its usage. The single model entry point.""" model = await self._resolve_model(employee) - credential = await self.credentials.get_active_for_provider(model.provider_id) + # Pool-aware: rotation providers (agy) fan out across every usable + # account; key-based providers simply get their one active credential. + pool = await self.credentials.list_active_for_provider(model.provider_id) + credential = pool[0] if pool else None response = await self.provider.complete( - model=model, credential=credential, messages=messages, system=system + model=model, + credential=credential, + credentials=pool or None, + messages=messages, + system=system, ) await self._record_usage( diff --git a/apps/api/app/core/config.py b/apps/api/app/core/config.py index 68c6a50..703d531 100644 --- a/apps/api/app/core/config.py +++ b/apps/api/app/core/config.py @@ -64,6 +64,21 @@ class Settings(BaseSettings): openai_compatible_base_url: str = "" openai_compatible_api_key: str = "" + # --- agy (Google Antigravity) -------------------------------------------- + #: OAuth client of the public Antigravity CLI application. Overridable so a + #: private Google Cloud OAuth client can be substituted without a code change. + agy_client_id: str = ( + "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com" + ) + agy_client_secret: str = "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" + #: Must stay loopback — Google only allows it for public installed clients. + agy_redirect_uri: str = "http://localhost:51121/oauth-callback" + agy_user_agent: str = "antigravity/1.11.5 windows/amd64" + #: Seconds an account is benched after Google answers 429 before the retry + #: that reverses the cooldown succeeds. Conservative because the real reset + #: wall is only known from quota headers when present. + agy_429_cooldown_seconds: int = 900 + # --- git ----------------------------------------------------------------- git_provider: str = "mock" gitea_url: str = "" diff --git a/apps/api/app/domain/ai_control/models.py b/apps/api/app/domain/ai_control/models.py index 826c800..9088b5f 100644 --- a/apps/api/app/domain/ai_control/models.py +++ b/apps/api/app/domain/ai_control/models.py @@ -19,6 +19,19 @@ class ProviderType(StrEnum): ANTHROPIC = "anthropic" GEMINI = "gemini" OPENROUTER = "openrouter" + #: Google Antigravity (``agy``) — OAuth accounts billed through a Google AI + #: subscription. Credentials of this type are *accounts*, not API keys. + AGY = "agy" + + +class CredentialStatus(StrEnum): + """Lifecycle of one agy account in the rotation pool.""" + + ACTIVE = "active" + #: Rate limited by Google; skipped until ``cooldown_until``. + COOLDOWN = "cooldown" + #: Refresh token rejected or another hard error; needs re-login. + ERROR = "error" class RoutingMode(StrEnum): @@ -61,12 +74,29 @@ class Provider: @dataclass(slots=True, kw_only=True) class ModelCredential: + """One credential for a provider. + + For key-based providers ``encrypted_api_key`` holds the encrypted API key and + the agy fields stay empty. For the agy provider the credential *is an + account*: ``encrypted_api_key`` holds the encrypted JSON blob + ``{"refresh_token": ..., "project_id": ..., "email": ...}`` while ``email`` / + ``project_id`` duplicate the identity columns so the admin console can show + and rotate accounts without touching the ciphertext. + """ + provider_id: str name: str encrypted_api_key: str id: str = field(default_factory=new_id) is_active: bool = True created_at: datetime = field(default_factory=utcnow) + #: agy account identity (empty for API-key credentials). + email: str = "" + project_id: str = "" + status: str = CredentialStatus.ACTIVE.value + cooldown_until: datetime | None = None + last_error: str = "" + last_used_at: datetime | None = None @dataclass(slots=True, kw_only=True) diff --git a/apps/api/app/infrastructure/adapters/agy.py b/apps/api/app/infrastructure/adapters/agy.py new file mode 100644 index 0000000..ee61b13 --- /dev/null +++ b/apps/api/app/infrastructure/adapters/agy.py @@ -0,0 +1,904 @@ +"""agy provider adapter — Google Antigravity accounts behind ModelProviderPort. + +``agy`` is Google's Antigravity: a Google-AI-subscription-billed route to the +Gemini 3 / Claude 4.5 / GPT-OSS families. There is no public API key; instead +each *account* is a Google OAuth identity whose refresh token is exchanged for +short-lived access tokens. This module implements that protocol directly, so +the platform never shells out to the CLI binary: + +* :class:`AgyAuthClient` — PKCE OAuth in two shapes: **listener mode** (a + loopback HTTP server captures the redirect; used when the operator's browser + runs on the same machine) and **paste mode** (the operator copies the + callback URL from the address bar and hands it back; the headless/remote + case — the page after redirect never finishes loading, and that is expected). +* :class:`AgyModelAdapter` — the ``ModelProviderPort`` implementation. Converts + our message format to the ``/v1internal:generateContent`` envelope, calls + Google's cloudcode-pa backend, parses the Gemini-shaped response. +* :class:`AgyAccountPool` state inside the adapter — round-robin across the + logged-in accounts with automatic 429 benching: a rate-limited account is + skipped until its cooldown expires, both in-process (fast) and in the + database (so background workers respect a cooldown set by the API). + +Nothing here ever returns or logs a refresh token: :func:`load_agy_accounts` +hands the adapter decrypted material in memory only, and admin endpoints see +email/project/status, never ciphertext. +""" + +from __future__ import annotations + +import asyncio +import base64 +import hashlib +import json +import logging +import secrets +import time +import uuid +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, field +from datetime import UTC, datetime, timedelta +from urllib.parse import parse_qs, urlencode, urlparse + +import httpx + +from app.application.ports import ChatMessage, ModelHealth, ModelResponse +from app.core.config import settings +from app.core.crypto import decrypt +from app.core.errors import ProviderUnavailableError, ValidationError +from app.domain.ai_control.models import AiModel, ModelCredential +from app.domain.base import utcnow + +logger = logging.getLogger("agy.adapter") + +# --------------------------------------------------------------------------- +# Protocol constants (the public Antigravity CLI OAuth application and the +# cloudcode-pa backend it talks to). +# --------------------------------------------------------------------------- + +OAUTH_AUTH_URL = "https://accounts.google.com/o/oauth2/v2/auth" +OAUTH_TOKEN_URL = "https://oauth2.googleapis.com/token" +USERINFO_URL = "https://www.googleapis.com/oauth2/v1/userinfo?alt=json" + +#: ``daily-cloudcode-pa`` first: it carries the free daily quota and 429s more +#: gracefully. ``cloudcode-pa`` is the fallback. +BASE_URLS: tuple[str, ...] = ( + "https://daily-cloudcode-pa.sandbox.googleapis.com", + "https://cloudcode-pa.googleapis.com", +) +PATH_GENERATE = "/v1internal:generateContent" +PATH_MODELS = "/v1internal:fetchAvailableModels" +PATH_LOAD_PROJECT = "/v1internal:loadCodeAssist" + +OAUTH_SCOPES: tuple[str, ...] = ( + "https://www.googleapis.com/auth/cloud-platform", + "https://www.googleapis.com/auth/userinfo.email", + "https://www.googleapis.com/auth/userinfo.profile", + "https://www.googleapis.com/auth/cclog", + "https://www.googleapis.com/auth/experimentsandconfigs", +) + +#: Same all-OFF matrix the Antigravity client sends. Without it Google applies +#: its defaults and refuses coding prompts it considers dangerous. +SAFETY_SETTINGS: list[dict[str, str]] = [ + {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "OFF"}, + {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "OFF"}, + {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "OFF"}, + {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "OFF"}, + {"category": "HARM_CATEGORY_CIVIC_INTEGRITY", "threshold": "OFF"}, +] + +#: How long a started login stays redeemable. Two OAuth round-trips plus a +#: manual paste must fit inside it comfortably. +CHALLENGE_TTL_SECONDS = 600 + +_StatusFn = Callable[..., Awaitable[None]] +_LoaderFn = Callable[[], Awaitable[list[ModelCredential]]] + + +# --------------------------------------------------------------------------- +# Small helpers +# --------------------------------------------------------------------------- + + +def _aware(value: datetime | None) -> datetime | None: + """SQLite hands back naive datetimes; comparisons need aware ones.""" + if value is None: + return None + if value.tzinfo is None: + return value.replace(tzinfo=UTC) + return value + + +def _pkce_pair() -> tuple[str, str]: + verifier = base64.urlsafe_b64encode(secrets.token_bytes(32)).rstrip(b"=").decode() + digest = hashlib.sha256(verifier.encode()).digest() + challenge = base64.urlsafe_b64encode(digest).rstrip(b"=").decode() + return verifier, challenge + + +def _random_session_id() -> str: + return f"-{secrets.randbelow(9_000_000_000_000_000_000)}" + + +def _fallback_project_id() -> str: + """Ad-hoc project when ``loadCodeAssist`` did not report one.""" + adjectives = ("useful", "bright", "swift", "calm", "bold") + nouns = ("fuze", "wave", "spark", "flow", "core") + adj = secrets.choice(adjectives) + noun = secrets.choice(nouns) + return f"{adj}-{noun}-{uuid.uuid4().hex[:5].lower()}" + + +def _retry_after_seconds(value: str | None) -> int | None: + if not value: + return None + try: + seconds = int(value.strip()) + except ValueError: + return None + return max(1, min(seconds, 4 * 3600)) + + +# --------------------------------------------------------------------------- +# Internal failure taxonomy (never crosses the port boundary) +# --------------------------------------------------------------------------- + + +class _AgyRateLimited(Exception): + def __init__(self, retry_after: int | None) -> None: + super().__init__("agy account rate limited") + self.retry_after = retry_after + + +class _AgyAuthError(Exception): + """Refresh token rejected / account not entitled — needs a fresh login.""" + + +class _AgyTransientError(Exception): + """5xx or transport failure on one backend URL — try the next account.""" + + +# --------------------------------------------------------------------------- +# OAuth login +# --------------------------------------------------------------------------- + + +@dataclass(slots=True, kw_only=True) +class AgyLoginChallenge: + login_id: str + auth_url: str + redirect_uri: str + verifier: str + state: str + created_at: float + + +@dataclass(slots=True, kw_only=True) +class AgyLoginResult: + """Decrypted account material, produced by a completed OAuth exchange.""" + + email: str + project_id: str + refresh_token: str + + +class AgyOAuthError(ValidationError): + """The pasted URL or the token exchange was rejected by Google.""" + + +class AgyAuthClient: + """Google OAuth (PKCE, installed-app client) for agy accounts.""" + + def __init__( + self, + *, + client_id: str | None = None, + client_secret: str | None = None, + redirect_uri: str | None = None, + timeout: float = 30.0, + ) -> None: + self.client_id = client_id or settings.agy_client_id + self.client_secret = client_secret or settings.agy_client_secret + self.redirect_uri = redirect_uri or settings.agy_redirect_uri + self.timeout = timeout + + # -- flow ------------------------------------------------------------- + + def begin(self) -> AgyLoginChallenge: + """Open a login: returns the URL the operator must open in a browser. + + The verifier stays in the returned challenge — it is the proof that + whoever redeems the login is whoever started it, and it is never + serialised out of the process. + """ + verifier, challenge = _pkce_pair() + login_id = uuid.uuid4().hex + state = uuid.uuid4().hex + params = urlencode( + { + "client_id": self.client_id, + "response_type": "code", + "redirect_uri": self.redirect_uri, + "scope": " ".join(OAUTH_SCOPES), + "code_challenge": challenge, + "code_challenge_method": "S256", + "state": state, + "access_type": "offline", + "prompt": "consent", + } + ) + return AgyLoginChallenge( + login_id=login_id, + auth_url=f"{OAUTH_AUTH_URL}?{params}", + redirect_uri=self.redirect_uri, + verifier=verifier, + state=state, + created_at=time.monotonic(), + ) + + async def exchange(self, challenge: AgyLoginChallenge, callback_url: str) -> AgyLoginResult: + """Redeem a started login with the callback URL from the browser. + + The operator may paste the full ``http://localhost:51121/oauth-callback? + …`` URL (it never finishes loading — that is the documented shape of + this flow, the parameters are already in the address bar). + """ + parsed = urlparse(callback_url.strip()) + query = parse_qs(parsed.query) + if query.get("error"): + raise AgyOAuthError(f"Google خطای احراز هویت داد: {query['error'][0]}") + codes = query.get("code") or [] + states = query.get("state") or [] + if not codes or not states: + raise AgyOAuthError( + "URL بازگشت معتبر نیست؛ باید پارامترهای code و state را داشته باشد." + ) + if states[0] != challenge.state: + raise AgyOAuthError("state بررسی نشد؛ این URL به این درخواست لاگین تعلق ندارد.") + return await self._exchange_code(challenge, codes[0]) + + async def _exchange_code(self, challenge: AgyLoginChallenge, code: str) -> AgyLoginResult: + tokens = await self._post_token( + { + "client_id": self.client_id, + "client_secret": self.client_secret, + "code": code, + "grant_type": "authorization_code", + "redirect_uri": challenge.redirect_uri, + "code_verifier": challenge.verifier, + } + ) + refresh_token = tokens.get("refresh_token") or "" + if not refresh_token: + # Happens when the consent screen was already granted for these + # scopes: Google omits the refresh token. A re-run with + # ``prompt=consent`` (which begin() always sends) fixes it, so the + # operator message says exactly that. + raise AgyOAuthError( + "Google refresh_token برنگرداند. یک بار دیگر لاگین کنید و در صفحهٔ " + "Google حتماً «Block/Allow» را تأیید کنید." + ) + email, project_id = await asyncio.gather( + self._fetch_email(tokens.get("access_token", "")), + self._fetch_project_id(tokens.get("access_token", "")), + ) + return AgyLoginResult( + email=email, project_id=project_id, refresh_token=refresh_token + ) + + async def refresh(self, refresh_token: str) -> str: + tokens = await self._post_token( + { + "client_id": self.client_id, + "client_secret": self.client_secret, + "refresh_token": refresh_token, + "grant_type": "refresh_token", + } + ) + access_token = tokens.get("access_token") + if not access_token: + raise _AgyAuthError("Google refresh را نپذیرفت (توکن منقضی یا باطل شده).") + return access_token + + # -- Google endpoints -------------------------------------------------- + + async def _post_token(self, data: dict[str, str]) -> dict: + try: + async with httpx.AsyncClient(timeout=self.timeout) as client: + response = await client.post( + OAUTH_TOKEN_URL, + data=data, + headers={"Content-Type": "application/x-www-form-urlencoded"}, + ) + except httpx.HTTPError as exc: + raise ProviderUnavailableError("اتصال به سرور OAuth گوگل برقرار نشد.") from exc + if response.status_code >= 400: + body = response.text[:400] + # A rejected refresh token is an account problem, not an outage. + if "invalid_grant" in body: + raise _AgyAuthError(f"Google refresh را نپذیرفت: {body}") + raise AgyOAuthError(f"تبادل توکن ناموفق بود ({response.status_code}): {body}") + return response.json() + + async def _fetch_email(self, access_token: str) -> str: + try: + async with httpx.AsyncClient(timeout=self.timeout) as client: + response = await client.get( + USERINFO_URL, headers={"Authorization": f"Bearer {access_token}"} + ) + if response.is_success: + return str(response.json().get("email", "")) + except httpx.HTTPError: + logger.debug("agy userinfo fetch failed", exc_info=True) + return "" + + async def _fetch_project_id(self, access_token: str) -> str: + for base in BASE_URLS: + try: + async with httpx.AsyncClient(timeout=self.timeout) as client: + response = await client.post( + f"{base}{PATH_LOAD_PROJECT}", + headers=self._api_headers(access_token), + json={ + "metadata": { + "ideType": "IDE_UNSPECIFIED", + "platform": "PLATFORM_UNSPECIFIED", + "pluginType": "GEMINI", + } + }, + ) + if response.is_success: + project = response.json().get("cloudaicompanionProject") + if isinstance(project, str) and project: + return project + if isinstance(project, dict) and project.get("id"): + return str(project["id"]) + except httpx.HTTPError: + continue + return "" + + def _api_headers(self, access_token: str) -> dict[str, str]: + return { + "Authorization": f"Bearer {access_token}", + "Content-Type": "application/json", + "User-Agent": settings.agy_user_agent, + } + + +# --------------------------------------------------------------------------- +# Pending-login store (process-local, bounded, TTL'd) +# --------------------------------------------------------------------------- + +_CHALLENGES: dict[str, AgyLoginChallenge] = {} +_CHALLENGES_LOCK = asyncio.Lock() + + +def _prune_challenges(now: float) -> None: + expired = [ + login_id + for login_id, challenge in _CHALLENGES.items() + if now - challenge.created_at > CHALLENGE_TTL_SECONDS + ] + for login_id in expired: + _CHALLENGES.pop(login_id, None) + + +async def begin_login() -> tuple[str, str]: + """Start a headless login. Returns ``(login_id, auth_url)``. + + ``auth_url`` is what the admin console shows the platform owner; the owner + opens it, signs in with Google, and the resulting callback URL (or, in + listener mode, the captured redirect) is redeemed against ``login_id``. + """ + challenge = AgyAuthClient().begin() + now = time.monotonic() + async with _CHALLENGES_LOCK: + _prune_challenges(now) + if len(_CHALLENGES) >= 32: + # Bound memory: pending logins are transient by nature. + oldest = min(_CHALLENGES, key=lambda k: _CHALLENGES[k].created_at) + _CHALLENGES.pop(oldest, None) + _CHALLENGES[challenge.login_id] = challenge + return challenge.login_id, challenge.auth_url + + +async def redeem_login(login_id: str, callback_url: str) -> AgyLoginResult: + challenge = _CHALLENGES.pop(login_id, None) + if challenge is None: + raise AgyOAuthError("این درخواست لاگین منقضی شده است؛ دوباره آغاز کنید.") + return await AgyAuthClient().exchange(challenge, callback_url) + + +async def redeem_login_by_waiting( + login_id: str, *, timeout_seconds: float = 300 +) -> AgyLoginResult: + """Listener mode: capture the redirect on the loopback callback port. + + Used by the operator-side script (and by a same-machine deployment): Google + redirects the browser to ``redirect_uri``; this server catches it, closes, + and redeems the login without any copy-paste. + """ + challenge = _CHALLENGES.get(login_id) + if challenge is None: + raise AgyOAuthError("این درخواست لاگین منقضی شده است؛ دوباره آغاز کنید.") + callback_path = urlparse(challenge.redirect_uri).path or "/" + port = int(urlparse(challenge.redirect_uri).port or 51121) + loop = asyncio.get_running_loop() + future: asyncio.Future[str] = loop.create_future() + + async def _handle( + reader: asyncio.StreamReader, writer: asyncio.StreamWriter + ) -> None: + try: + request_line = await reader.readline() + parts = request_line.decode("latin-1").split(" ") + target = parts[1] if len(parts) > 1 else "/" + if urlparse(target).path != callback_path: + writer.write(b"HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\n\r\n") + await writer.drain() + return + writer.write( + b"HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=utf-8\r\n" + b"Content-Length: 46\r\nConnection: close\r\n\r\n" + b"OK - you may close this tab." + ) + await writer.drain() + if not future.done(): + future.set_result(target) + except Exception: # noqa: BLE001 - the future carries the real error + if not future.done(): + future.set_exception(AgyOAuthError("دریافت بازگشت OAuth ناموفق بود.")) + finally: + writer.close() + + server = await asyncio.start_server(_handle, "127.0.0.1", port) + try: + async with asyncio.timeout(timeout_seconds): + target = await future + except TimeoutError as exc: + raise AgyOAuthError("مهلت لاگین تمام شد؛ دوباره آغاز کنید.") from exc + finally: + server.close() + await server.wait_closed() + _CHALLENGES.pop(login_id, None) + redirect_uri = challenge.redirect_uri + return await AgyAuthClient(redirect_uri=redirect_uri).exchange( + challenge, f"http://127.0.0.1:{port}{target}" + ) + + +# --------------------------------------------------------------------------- +# Account pool +# --------------------------------------------------------------------------- + + +@dataclass(slots=True, kw_only=True) +class AgyAccount: + """One usable account: decrypted material plus its rotation identity.""" + + credential_id: str + email: str + project_id: str + refresh_token: str + + +@dataclass(slots=True) +class _PoolState: + #: Monotonic rotation cursor — spreads load across the account set. + cursor: int = 0 + #: credential_id -> epoch seconds until which the account is benched. + benched_until: dict[str, float] = field(default_factory=dict) + #: credential_id -> (access token, expiry epoch) — avoids a Google + #: round-trip on every single completion. + tokens: dict[str, tuple[str, float]] = field(default_factory=dict) + + +def _account_of(credential: ModelCredential) -> AgyAccount: + """Decrypt one credential row into usable login material. + + The ciphertext is a JSON envelope written at login time; older key-shaped + credentials (empty email) are treated as unusable for agy. + """ + try: + payload = json.loads(decrypt(credential.encrypted_api_key)) + return AgyAccount( + credential_id=credential.id, + email=str(payload.get("email") or credential.email or ""), + project_id=str(payload.get("project_id") or credential.project_id or ""), + refresh_token=str(payload.get("refresh_token") or ""), + ) + except Exception: + return AgyAccount( + credential_id=credential.id, + email=credential.email, + project_id=credential.project_id, + refresh_token="", + ) + + +class AgyModelAdapter: + """ModelProviderPort over the agy account pool. + + One instance per process (the adapter registry caches it), so the rotation + cursor and the in-memory bench are shared by every worker and every + request. Hard state (cooldown / error / last-used) is mirrored to the + database through the injected ``report`` callback so the background + workers see the same picture. + """ + + def __init__( + self, + *, + loader: _LoaderFn | None = None, + report: _StatusFn | None = None, + timeout: float = 240.0, + ) -> None: + self._loader = loader + self._report = report + self._timeout = timeout + self._state = _PoolState() + self._lock = asyncio.Lock() + + # -- port --------------------------------------------------------------- + + async def complete( + self, + *, + model: AiModel, + credential: ModelCredential | None = None, + credentials: list[ModelCredential] | None = None, + messages: list[ChatMessage], + system: str = "", + temperature: float = 0.2, + max_tokens: int = 2048, + ) -> ModelResponse: + pool = credentials if credentials is not None else ( + [credential] if credential is not None else await self._load() + ) + accounts = [a for a in (_account_of(c) for c in pool) if a.refresh_token] + if not accounts: + raise ProviderUnavailableError( + "هیچ اکانت agy فعالی ثبت نشده است. مالک سامانه باید ابتدا از " + "کنسول مدیریت یک اکانت Google اضافه کند." + ) + ordered = self._order_accounts(accounts) + failures: list[str] = [] + for account in ordered: + try: + response = await self._complete_with( + account, model=model, messages=messages, system=system, + temperature=temperature, max_tokens=max_tokens, + ) + except _AgyRateLimited as exc: + failures.append(f"{account.email or account.credential_id}: محدودیت نرخ") + await self._bench(account, exc.retry_after) + continue + except _AgyAuthError as exc: + failures.append(f"{account.email or account.credential_id}: {exc}") + await self._bench(account, 3600, status="error", error=str(exc)) + continue + except _AgyTransientError as exc: + failures.append(f"{account.email or account.credential_id}: {exc}") + continue + await self._succeed(account) + return response + raise ProviderUnavailableError( + "همهٔ اکانت‌های agy در دسترس نبودند: " + "; ".join(failures[:4]) + ) + + async def health(self) -> ModelHealth: + accounts = [a for a in (_account_of(c) for c in await self._load()) if a.refresh_token] + if not accounts: + return ModelHealth(ok=False, detail="هیچ اکانت agy ثبت نشده است.") + try: + await self._access_token(accounts[0]) + except _AgyAuthError as exc: + return ModelHealth(ok=False, detail=str(exc)) + except Exception as exc: # noqa: BLE001 - health must not raise + return ModelHealth(ok=False, detail=str(exc)) + return ModelHealth(ok=True, detail=f"{len(accounts)} اکانت فعال") + + async def available_models(self) -> list[str]: + """Models the logged-in subscription exposes (admin console helper).""" + accounts = [a for a in (_account_of(c) for c in await self._load()) if a.refresh_token] + if not accounts: + return [] + token = await self._access_token(accounts[0]) + for base in BASE_URLS: + try: + async with httpx.AsyncClient(timeout=30.0) as client: + response = await client.post( + f"{base}{PATH_MODELS}", + headers={ + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + "User-Agent": settings.agy_user_agent, + }, + json={}, + ) + if response.is_success: + models = response.json().get("models") or {} + return sorted(models) + except httpx.HTTPError: + continue + return [] + + # -- internals ---------------------------------------------------------- + + async def _load(self) -> list[ModelCredential]: + if self._loader is None: + return [] + return await self._loader() + + async def _report_state( + self, + credential_id: str, + *, + status: str | None = None, + cooldown_until: datetime | None = None, + error: str | None = None, + used: bool = False, + ) -> None: + if self._report is None: + return + try: + await self._report( + credential_id, + status=status, + cooldown_until=cooldown_until, + error=error, + used=used, + ) + except Exception: # noqa: BLE001 - reporting must never kill a call + logger.debug("agy state report failed", exc_info=True) + + def _order_accounts(self, accounts: list[AgyAccount]) -> list[AgyAccount]: + """Round-robin, skipping accounts benched in this process. + + The DB-shaped cooldown is already applied by the query that produced + the pool; this pass additionally honours recent 429s that have not been + persisted yet (or a reporter that is not wired). + """ + now = time.time() + usable = [a for a in accounts if self._state.benched_until.get(a.credential_id, 0) <= now] + if not usable: + # Everything is benched in memory — fall through to the DB's view + # rather than failing a request the pool could still serve. + usable = list(accounts) + n = len(usable) + self._state.cursor += 1 + start = self._state.cursor % n + return usable[start:] + usable[:start] + + async def _bench( + self, + account: AgyAccount, + retry_after: int | None, + *, + status: str = "cooldown", + error: str = "", + ) -> None: + seconds = retry_after or settings.agy_429_cooldown_seconds + until = datetime.now(UTC) + timedelta(seconds=seconds) + async with self._lock: + self._state.benched_until[account.credential_id] = until.timestamp() + await self._report_state( + account.credential_id, status=status, cooldown_until=until, error=error + ) + + async def _succeed(self, account: AgyAccount) -> None: + async with self._lock: + self._state.benched_until.pop(account.credential_id, None) + await self._report_state(account.credential_id, status="active", used=True) + + async def _access_token(self, account: AgyAccount) -> str: + """Cached refresh: Google tokens live ~1h; one call serves the burst.""" + cached = self._state.tokens.get(account.credential_id) + if cached and cached[1] > time.time() + 30: + return cached[0] + token = await AgyAuthClient().refresh(account.refresh_token) + async with self._lock: + self._state.tokens[account.credential_id] = (token, time.time() + 3300) + return token + + async def _complete_with( + self, + account: AgyAccount, + *, + model: AiModel, + messages: list[ChatMessage], + system: str, + temperature: float, + max_tokens: int, + ) -> ModelResponse: + token = await self._access_token(account) + payload = _build_generate_payload( + account, model.internal_name, messages, system, temperature, max_tokens + ) + headers = { + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + "User-Agent": settings.agy_user_agent, + "Accept": "application/json", + } + started = time.perf_counter() + last_transient: Exception | None = None + retry_after: int | None = None + for base in BASE_URLS: + try: + async with httpx.AsyncClient(timeout=self._timeout) as client: + response = await client.post( + f"{base}{PATH_GENERATE}", json=payload, headers=headers + ) + except httpx.HTTPError: + last_transient = _AgyTransientError(f"اتصال به {base} برقرار نشد.") + continue + if response.status_code == 429: + # Both backends share one quota, but a 429 can also be a + # backend hiccup — try the other base URL before benching the + # whole account. + retry_after = _retry_after_seconds( + response.headers.get("Retry-After") + ) or retry_after + continue + if response.status_code in (401, 403): + raise _AgyAuthError( + f"اکانت {account.email or 'agy'} دسترسی ندارد یا نشست منقضی شده است." + ) + if response.status_code >= 500: + last_transient = _AgyTransientError(f"{base} خطای {response.status_code} داد.") + continue + if response.status_code >= 400: + # Our request shape is wrong — another account will fail the + # same way, so surface it instead of burning the pool. + raise ProviderUnavailableError( + f"agy درخواست را رد کرد ({response.status_code}): " + f"{response.text[:300]}" + ) + return _parse_generate_response( + response.json(), model=model, started=started + ) + if retry_after is not None: + raise _AgyRateLimited(retry_after) + raise last_transient or _AgyTransientError("هیچ بک‌اند agy پاسخ نداد.") + + +# --------------------------------------------------------------------------- +# Message / payload translation (portable, pure functions — unit-tested) +# --------------------------------------------------------------------------- + + +def _build_generate_payload( + account: AgyAccount, + model_name: str, + messages: list[ChatMessage], + system: str, + temperature: float, + max_tokens: int, +) -> dict: + """Translate the port's OpenAI-shaped call into the Antigravity envelope.""" + system_parts: list[str] = [system] if system else [] + contents: list[dict] = [] + for message in messages: + if message.role == "system": + system_parts.append(message.content) + continue + role = "model" if message.role == "assistant" else "user" + contents.append({"role": role, "parts": [{"text": message.content}]}) + + request: dict = { + "contents": contents, + "sessionId": _random_session_id(), + "generationConfig": { + "temperature": temperature, + "maxOutputTokens": max_tokens, + }, + "toolConfig": {"functionCallingConfig": {"mode": "VALIDATED"}}, + "safetySettings": SAFETY_SETTINGS, + } + if system_parts: + request["systemInstruction"] = { + "role": "user", + "parts": [{"text": "\n\n".join(p for p in system_parts if p)}], + } + return { + "project": account.project_id or _fallback_project_id(), + "model": model_name, + "userAgent": "antigravity", + "requestId": f"agent-{uuid.uuid4()}", + "request": request, + } + + +def _parse_generate_response( + data: dict, *, model: AiModel, started: float +) -> ModelResponse: + candidates = data.get("candidates") or [{}] + candidate = candidates[0] if isinstance(candidates[0], dict) else {} + text_parts: list[str] = [] + for part in (candidate.get("content") or {}).get("parts") or []: + if not isinstance(part, dict): + continue + text = part.get("text") + if text and not part.get("thought"): + text_parts.append(text) + usage = data.get("usageMetadata") or {} + latency_ms = int((time.perf_counter() - started) * 1000) + return ModelResponse( + content="".join(text_parts), + model_internal_name=model.internal_name, + input_tokens=int(usage.get("promptTokenCount", 0) or 0), + output_tokens=int(usage.get("candidatesTokenCount", 0) or 0), + latency_ms=latency_ms, + raw={ + "finish_reason": candidate.get("finishReason", ""), + "thought_tokens": int(usage.get("thoughtsTokenCount", 0) or 0), + }, + ) + + +# --------------------------------------------------------------------------- +# Database wiring for the registry (each call opens its own short session, so +# the adapter never outlives or depends on an HTTP request's session). +# --------------------------------------------------------------------------- + + +async def load_agy_accounts() -> list[ModelCredential]: + """Every usable agy credential row, ordered least-recently-used first.""" + from app.core.database import SessionFactory + from app.infrastructure.repositories.ai_control import ( + ModelCredentialRepository, + ProviderRepository, + ) + + async with SessionFactory() as session: + providers = await ProviderRepository(session).list() + provider = next( + (p for p in providers if p.type.value == "agy" and p.is_active), None + ) + if provider is None: + return [] + return await ModelCredentialRepository(session).list_active_for_provider( + provider.id + ) + + +async def report_agy_state( + credential_id: str, + *, + status: str | None = None, + cooldown_until: datetime | None = None, + error: str | None = None, + used: bool = False, +) -> None: + """Mirror rotation state into the DB for the workers' benefit.""" + from app.core.database import SessionFactory + from app.infrastructure.repositories.ai_control import ModelCredentialRepository + + async with SessionFactory() as session: + repo = ModelCredentialRepository(session) + if status is not None: + await repo.set_status( + credential_id, + status=status, + cooldown_until=cooldown_until, + error=error or "", + ) + if used: + await repo.mark_used(credential_id, when=utcnow()) + await session.commit() + + +__all__ = [ + "AgyAccount", + "AgyAuthClient", + "AgyLoginChallenge", + "AgyLoginResult", + "AgyModelAdapter", + "AgyOAuthError", + "begin_login", + "load_agy_accounts", + "redeem_login", + "redeem_login_by_waiting", + "report_agy_state", +] diff --git a/apps/api/app/infrastructure/adapters/mock.py b/apps/api/app/infrastructure/adapters/mock.py index a81479b..263cac4 100644 --- a/apps/api/app/infrastructure/adapters/mock.py +++ b/apps/api/app/infrastructure/adapters/mock.py @@ -66,6 +66,7 @@ class MockModelProviderAdapter: *, model: AiModel, credential=None, + credentials=None, # accepted for port parity; the mock has no pool messages: list[ChatMessage], system: str = "", temperature: float = 0.2, diff --git a/apps/api/app/infrastructure/adapters/openai_compatible.py b/apps/api/app/infrastructure/adapters/openai_compatible.py index d3d4620..c587ddc 100644 --- a/apps/api/app/infrastructure/adapters/openai_compatible.py +++ b/apps/api/app/infrastructure/adapters/openai_compatible.py @@ -34,7 +34,8 @@ class OpenAICompatibleModelAdapter: self, *, model: AiModel, - credential: ModelCredential | None, + credential: ModelCredential | None = None, + credentials: list[ModelCredential] | None = None, # no pool here messages: list[ChatMessage], system: str = "", temperature: float = 0.2, diff --git a/apps/api/app/infrastructure/adapters/registry.py b/apps/api/app/infrastructure/adapters/registry.py index 641b0b8..65b28bc 100644 --- a/apps/api/app/infrastructure/adapters/registry.py +++ b/apps/api/app/infrastructure/adapters/registry.py @@ -28,9 +28,17 @@ def get_model_provider(): base_url=settings.openai_compatible_base_url, default_api_key=settings.openai_compatible_api_key, ) + if provider == "agy": + from app.infrastructure.adapters.agy import ( + AgyModelAdapter, + load_agy_accounts, + report_agy_state, + ) + + return AgyModelAdapter(loader=load_agy_accounts, report=report_agy_state) raise ProviderUnavailableError( f"AI_DEFAULT_PROVIDER='{provider}' پشتیبانی نمی‌شود. " - "مقادیر مجاز: mock | openai_compatible" + "مقادیر مجاز: mock | openai_compatible | agy" ) diff --git a/apps/api/app/infrastructure/db/models.py b/apps/api/app/infrastructure/db/models.py index 17111ac..410ce35 100644 --- a/apps/api/app/infrastructure/db/models.py +++ b/apps/api/app/infrastructure/db/models.py @@ -132,6 +132,17 @@ class AiModelCredentialRow(Base): encrypted_api_key: Mapped[str] = mapped_column(Text, default="", nullable=False) is_active: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False) created_at: Mapped[datetime] = _created() + # --- agy account columns (empty for API-key credentials) ----------------- + #: Google account email this credential belongs to; also the rotation label. + email: Mapped[str] = mapped_column(String(320), default="", nullable=False) + #: Antigravity ``cloudaicompanionProject`` id resolved during login. + project_id: Mapped[str] = mapped_column(String(160), default="", nullable=False) + #: ``active`` | ``cooldown`` | ``error`` — see ``CredentialStatus``. + status: Mapped[str] = mapped_column(String(16), default="active", nullable=False) + #: When set and in the future, the rotation pool skips this account. + cooldown_until: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True) + last_error: Mapped[str] = mapped_column(String(512), default="", nullable=False) + last_used_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True) class AiModelRow(Base): diff --git a/apps/api/app/infrastructure/repositories/ai_control.py b/apps/api/app/infrastructure/repositories/ai_control.py index 4b84a37..7afa60d 100644 --- a/apps/api/app/infrastructure/repositories/ai_control.py +++ b/apps/api/app/infrastructure/repositories/ai_control.py @@ -7,6 +7,8 @@ move rows. from __future__ import annotations +from datetime import datetime + from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession @@ -20,6 +22,7 @@ from app.domain.ai_control.models import ( RoutingMode, RoutingPolicy, ) +from app.domain.base import utcnow from app.infrastructure.db.models import ( AiModelCredentialRow, AiModelRow, @@ -41,6 +44,9 @@ def _credential_to_domain(row: AiModelCredentialRow) -> ModelCredential: id=row.id, provider_id=row.provider_id, name=row.name, encrypted_api_key=row.encrypted_api_key, is_active=row.is_active, created_at=row.created_at, + email=row.email or "", project_id=row.project_id or "", + status=row.status or "active", cooldown_until=row.cooldown_until, + last_error=row.last_error or "", last_used_at=row.last_used_at, ) @@ -128,22 +134,115 @@ class ModelCredentialRepository: .where( AiModelCredentialRow.provider_id == provider_id, AiModelCredentialRow.is_active.is_(True), + AiModelCredentialRow.status == "active", ) .order_by(AiModelCredentialRow.created_at.desc()) ) row = (await self.session.execute(stmt)).scalars().first() return _credential_to_domain(row) if row else None + async def list_active_for_provider(self, provider_id: str) -> list[ModelCredential]: + """Every credential the rotation pool may use. + + An account in ``cooldown`` whose window has passed is usable again, so + the status filter is applied in SQL only for hard errors; the pool + itself checks ``cooldown_until`` against the clock to stay race-free + with concurrent requests. + """ + now = utcnow() + stmt = ( + select(AiModelCredentialRow) + .where( + AiModelCredentialRow.provider_id == provider_id, + AiModelCredentialRow.is_active.is_(True), + AiModelCredentialRow.status != "error", + (AiModelCredentialRow.status != "cooldown") + | (AiModelCredentialRow.cooldown_until.is_(None)) + | (AiModelCredentialRow.cooldown_until <= now), + ) + .order_by(AiModelCredentialRow.last_used_at.asc().nulls_first()) + ) + return [ + _credential_to_domain(r) + for r in (await self.session.execute(stmt)).scalars() + ] + + async def find_by_email(self, provider_id: str, email: str) -> ModelCredential | None: + """One account per Google address — re-login replaces, never duplicates.""" + stmt = select(AiModelCredentialRow).where( + AiModelCredentialRow.provider_id == provider_id, + AiModelCredentialRow.email == email, + ) + row = (await self.session.execute(stmt)).scalars().first() + return _credential_to_domain(row) if row else None + async def create(self, credential: ModelCredential) -> ModelCredential: row = AiModelCredentialRow( id=credential.id, provider_id=credential.provider_id, name=credential.name, encrypted_api_key=credential.encrypted_api_key, is_active=credential.is_active, created_at=credential.created_at, + email=credential.email, project_id=credential.project_id, + status=credential.status, cooldown_until=credential.cooldown_until, + last_error=credential.last_error, last_used_at=credential.last_used_at, ) self.session.add(row) await self.session.flush() return _credential_to_domain(row) + async def replace_material( + self, credential_id: str, *, encrypted_api_key: str, project_id: str + ) -> ModelCredential | None: + """Re-login of a known account: swap ciphertext and clear error state.""" + row = await self.session.get(AiModelCredentialRow, credential_id) + if row is None: + return None + row.encrypted_api_key = encrypted_api_key + row.project_id = project_id + row.status = "active" + row.cooldown_until = None + row.last_error = "" + await self.session.flush() + return _credential_to_domain(row) + + async def set_status( + self, + credential_id: str, + *, + status: str, + cooldown_until: datetime | None = None, + error: str = "", + ) -> None: + row = await self.session.get(AiModelCredentialRow, credential_id) + if row is not None: + row.status = status + row.cooldown_until = cooldown_until + row.last_error = error + await self.session.flush() + + async def mark_used(self, credential_id: str, *, when: datetime) -> None: + row = await self.session.get(AiModelCredentialRow, credential_id) + if row is not None: + row.last_used_at = when + await self.session.flush() + + async def set_active(self, credential_id: str, active: bool) -> None: + row = await self.session.get(AiModelCredentialRow, credential_id) + if row is not None: + row.is_active = active + if active and row.status == "error": + row.status = "active" + row.cooldown_until = None + row.last_error = "" + await self.session.flush() + + async def delete(self, credential_id: str) -> bool: + row = await self.session.get(AiModelCredentialRow, credential_id) + if row is None: + return False + await self.session.delete(row) + await self.session.flush() + return True + class AiModelRepository: def __init__(self, session: AsyncSession) -> None: diff --git a/apps/api/app/interfaces/http/routers/agy.py b/apps/api/app/interfaces/http/routers/agy.py new file mode 100644 index 0000000..edf5fc0 --- /dev/null +++ b/apps/api/app/interfaces/http/routers/agy.py @@ -0,0 +1,343 @@ +"""agy account management — Google OAuth headless login for the AI provider. + +The agy provider has no API keys: an *account* is a Google identity. This +router implements the headless paste-URL login loop for the platform owner: + +1. ``POST /admin/agy/accounts/login/start`` — returns an ``auth_url`` the owner + opens in a browser and completes the Google consent screen. +2. Google redirects to the loopback callback URL. Nothing is listening there + on purpose — the parameters (``code``, ``state``) are already in the address + bar, and the owner pastes the whole URL back. +3. ``POST /admin/agy/accounts/login/complete`` — exchanges the code against + Google (PKCE proof stays server-side), fetches the account email and + Antigravity project id, and stores the refresh token **encrypted at rest**. + +Every response is deliberately free of credentials: the console sees email, +status and cooldown, never ciphertext. Same deny-by-default guard as the rest +of the admin console. +""" + +from __future__ import annotations + +import json +from datetime import datetime + +from fastapi import APIRouter, Depends, Request, status +from pydantic import BaseModel, ConfigDict, Field +from sqlalchemy.ext.asyncio import AsyncSession + +from app.application.services.audit_service import record_audit +from app.core.crypto import encrypt +from app.core.database import get_session +from app.core.errors import NotFoundError +from app.domain.ai_control.models import AiModel, ModelCredential, ProviderType +from app.domain.audit.models import AuditAction +from app.domain.identity.models import User +from app.infrastructure.repositories.ai_control import ( + ModelCredentialRepository, + ProviderRepository, +) +from app.interfaces.http.deps import client_meta, require_platform_owner + +router = APIRouter( + prefix="/admin/agy", + tags=["admin-agy"], + # Deny-by-default, same rule as the admin console: an endpoint added later + # is owner-only without its author having to remember. + dependencies=[Depends(require_platform_owner)], +) + +#: Name of the singleton agy provider row. The seed creates it; the login flow +#: creates it on demand so a bare deployment is never stuck. +AGY_PROVIDER_NAME = "agy" + +_PROJECTED_OUT_FIELDS = ( + "id", + "provider_id", + "name", + "email", + "project_id", + "status", + "cooldown_until", + "is_active", + "last_error", + "last_used_at", + "created_at", +) + + +# --------------------------------------------------------------------------- +# Schemas — note there is no field anywhere capable of holding a token. +# --------------------------------------------------------------------------- + + +class AgyLoginStartOut(BaseModel): + login_id: str + auth_url: str + + +class AgyLoginCompleteIn(BaseModel): + login_id: str = Field(min_length=8, max_length=64) + callback_url: str = Field(min_length=20, max_length=2048) + + +class AgyAccountOut(BaseModel): + model_config = ConfigDict(from_attributes=True) + + id: str + provider_id: str + name: str + email: str + project_id: str + status: str + cooldown_until: datetime | None + is_active: bool + last_error: str + last_used_at: datetime | None + created_at: datetime + + +class AgyTestOut(BaseModel): + ok: bool + detail: str + model_internal_name: str + input_tokens: int + output_tokens: int + + +# --------------------------------------------------------------------------- +# Provider helpers +# --------------------------------------------------------------------------- + + +async def _ensure_agy_provider(repo: ProviderRepository): + provider = await repo.get_by_name(AGY_PROVIDER_NAME) + if provider is not None: + return provider + from app.domain.ai_control.models import Provider + + return await repo.create(Provider(name=AGY_PROVIDER_NAME, type=ProviderType.AGY)) + + +def _account_out(credential: ModelCredential) -> AgyAccountOut: + return AgyAccountOut( + id=credential.id, + provider_id=credential.provider_id, + name=credential.name, + email=credential.email, + project_id=credential.project_id, + status=credential.status, + cooldown_until=credential.cooldown_until, + is_active=credential.is_active, + last_error=credential.last_error, + last_used_at=credential.last_used_at, + created_at=credential.created_at, + ) + + +# --------------------------------------------------------------------------- +# Endpoints +# --------------------------------------------------------------------------- + + +@router.post( + "/accounts/login/start", response_model=AgyLoginStartOut, + status_code=status.HTTP_201_CREATED, +) +async def start_agy_login() -> AgyLoginStartOut: + """Open a headless login and hand back the Google URL to open.""" + from app.infrastructure.adapters.agy import begin_login + + login_id, auth_url = await begin_login() + return AgyLoginStartOut(login_id=login_id, auth_url=auth_url) + + +@router.post( + "/accounts/login/complete", response_model=AgyAccountOut, + status_code=status.HTTP_201_CREATED, +) +async def complete_agy_login( + payload: AgyLoginCompleteIn, + request: Request, + owner: User = Depends(require_platform_owner), + session: AsyncSession = Depends(get_session), +) -> AgyAccountOut: + """Redeem a started login with the callback URL the owner pasted.""" + from app.core.errors import ValidationError + from app.infrastructure.adapters.agy import redeem_login + + try: + result = await redeem_login(payload.login_id, payload.callback_url) + except Exception as exc: + # Login failures are user-fixable input problems, never 500s. + raise ValidationError(str(exc)) from exc + if not result.email: + raise ValidationError("Google نشانی ایمیل اکانت را برنگرداند.") + + providers = ProviderRepository(session) + credentials = ModelCredentialRepository(session) + provider = await _ensure_agy_provider(providers) + + envelope = json.dumps( + { + "refresh_token": result.refresh_token, + "project_id": result.project_id, + "email": result.email, + } + ) + existing = await credentials.find_by_email(provider.id, result.email) + if existing is not None: + # Same Google account re-login: rotate the key, keep the row. + credential = await credentials.replace_material( + existing.id, encrypted_api_key=encrypt(envelope), + project_id=result.project_id, + ) + else: + credential = await credentials.create( + ModelCredential( + provider_id=provider.id, + name=result.email, + encrypted_api_key=encrypt(envelope), + email=result.email, + project_id=result.project_id, + ) + ) + assert credential is not None # narrowed for type checkers; create/replace return it + + ip, ua = client_meta(request) + await record_audit( + session, action=AuditAction.PROVIDER_CHANGE, actor_user_id=owner.id, + target_type="agy_account", target_id=credential.id, + metadata={"email": result.email, "provider_id": provider.id}, + ip=ip, user_agent=ua, + ) + return _account_out(credential) + + +@router.get("/accounts", response_model=list[AgyAccountOut]) +async def list_agy_accounts( + session: AsyncSession = Depends(get_session), +) -> list[AgyAccountOut]: + providers = ProviderRepository(session) + provider = await providers.get_by_name(AGY_PROVIDER_NAME) + if provider is None: + return [] + rows = await ModelCredentialRepository(session).list_for_provider(provider.id) + return [_account_out(c) for c in rows] + + +@router.post("/accounts/{credential_id}/deactivate", response_model=AgyAccountOut) +async def deactivate_agy_account( + credential_id: str, + request: Request, + owner: User = Depends(require_platform_owner), + session: AsyncSession = Depends(get_session), +) -> AgyAccountOut: + return await _set_account_active( + credential_id, active=False, request=request, owner=owner, session=session + ) + + +@router.post("/accounts/{credential_id}/activate", response_model=AgyAccountOut) +async def activate_agy_account( + credential_id: str, + request: Request, + owner: User = Depends(require_platform_owner), + session: AsyncSession = Depends(get_session), +) -> AgyAccountOut: + return await _set_account_active( + credential_id, active=True, request=request, owner=owner, session=session + ) + + +async def _set_account_active( + credential_id: str, *, active: bool, request: Request, owner: User, + session: AsyncSession, +) -> AgyAccountOut: + repo = ModelCredentialRepository(session) + credential = await repo.get(credential_id) + if credential is None: + raise NotFoundError("اکانت agy یافت نشد.") + await repo.set_active(credential_id, active) + credential = await repo.get(credential_id) + assert credential is not None + ip, ua = client_meta(request) + await record_audit( + session, action=AuditAction.PROVIDER_CHANGE, actor_user_id=owner.id, + target_type="agy_account", target_id=credential_id, + metadata={"active": active}, + ip=ip, user_agent=ua, + ) + return _account_out(credential) + + +@router.delete("/accounts/{credential_id}", status_code=status.HTTP_204_NO_CONTENT) +async def delete_agy_account( + credential_id: str, + request: Request, + owner: User = Depends(require_platform_owner), + session: AsyncSession = Depends(get_session), +) -> None: + repo = ModelCredentialRepository(session) + credential = await repo.get(credential_id) + if credential is None: + raise NotFoundError("اکانت agy یافت نشد.") + await repo.delete(credential_id) + ip, ua = client_meta(request) + await record_audit( + session, action=AuditAction.PROVIDER_CHANGE, actor_user_id=owner.id, + target_type="agy_account", target_id=credential_id, + metadata={"deleted": True, "email": credential.email}, + ip=ip, user_agent=ua, + ) + + +@router.post("/accounts/{credential_id}/test", response_model=AgyTestOut) +async def test_agy_account( + credential_id: str, + session: AsyncSession = Depends(get_session), +) -> AgyTestOut: + """Run a one-shot completion through exactly this account.""" + from app.application.ports import ChatMessage + from app.infrastructure.adapters.agy import AgyModelAdapter + + repo = ModelCredentialRepository(session) + credential = await repo.get(credential_id) + if credential is None: + raise NotFoundError("اکانت agy یافت نشد.") + + adapter = AgyModelAdapter() # stateless for a single forced account + model = AiModel( + provider_id=credential.provider_id, + internal_name="gemini-3-flash", + display_name_for_admin="تست اکانت", + ) + try: + response = await adapter.complete( + model=model, + credentials=[credential], + messages=[ChatMessage(role="user", content="سلام، یک کلمه پاسخ بده.")], + max_tokens=64, + ) + except Exception as exc: # noqa: BLE001 - the console wants the reason, not a 500 + return AgyTestOut( + ok=False, detail=str(exc)[:300], model_internal_name="", + input_tokens=0, output_tokens=0, + ) + return AgyTestOut( + ok=True, detail=response.content[:300], + model_internal_name=response.model_internal_name, + input_tokens=response.input_tokens, output_tokens=response.output_tokens, + ) + + +@router.get("/models", response_model=list[str]) +async def list_agy_models() -> list[str]: + """Model ids the logged-in agy subscription exposes (registry helper).""" + from app.infrastructure.adapters.registry import get_model_provider + + provider = get_model_provider() + lister = getattr(provider, "available_models", None) + if lister is None: + return [] + return await lister() diff --git a/apps/api/app/main.py b/apps/api/app/main.py index 83ee93c..94fe661 100644 --- a/apps/api/app/main.py +++ b/apps/api/app/main.py @@ -17,6 +17,7 @@ from app.interfaces.http.errors import register_exception_handlers from app.interfaces.http.request_context import RequestContextMiddleware from app.interfaces.http.routers import ( admin, + agy, auth, billing, health, @@ -121,6 +122,7 @@ def create_app() -> FastAPI: application.include_router(knowledge.router, prefix="/api") application.include_router(sandbox.router, prefix="/api") application.include_router(admin.router, prefix="/api") + application.include_router(agy.router, prefix="/api") @application.get("/", tags=["meta"]) async def root() -> dict: diff --git a/apps/api/app/seed.py b/apps/api/app/seed.py index 8f369fc..8f69bbc 100644 --- a/apps/api/app/seed.py +++ b/apps/api/app/seed.py @@ -2,7 +2,9 @@ Run with ``pnpm db:seed``. Creates a platform owner, a demo user, a sample organization that both of them own, the system agent templates, routing -policies, rate cards and a mock provider with two models. +policies, rate cards and a mock provider with two models. The agy (Google +Antigravity) provider and its model catalog are seeded too; agy models become +routable only when ``AI_DEFAULT_PROVIDER=agy``. Safe to run repeatedly: every step checks for existing rows first. @@ -18,6 +20,7 @@ import logging from sqlalchemy import select +from app.core.config import settings from app.core.database import Base, SessionFactory, engine from app.core.security import hash_password from app.domain.ai_control.models import ( @@ -60,6 +63,42 @@ DEMO_PASSWORD = "demo12345" DEMO_ORG_NAME = "سازمان نمونه" DEMO_ORG_SLUG = "sample-org" +#: The agy (Google Antigravity) model catalog. ``internal_name`` is the model id +#: the Antigravity backend expects; the display name is admin-only. Seeded for +#: every deployment, activated for routing only when agy is the default +#: provider — see the seed step that consumes this table. +AGY_MODELS: list[tuple[str, str, list[Capability]]] = [ + ( + "gemini-3-flash", + "Gemini 3 Flash", + [Capability.TEXT, Capability.CODE, Capability.LONG_CONTEXT], + ), + ( + "gemini-3-pro-high", + "Gemini 3 Pro", + [ + Capability.TEXT, Capability.REASONING, Capability.CODE, + Capability.VISION, Capability.LONG_CONTEXT, + ], + ), + ( + "claude-sonnet-4-5", + "Claude Sonnet 4.5", + [Capability.TEXT, Capability.CODE, Capability.TOOLS], + ), + ( + "claude-sonnet-4-5-thinking", + "Claude Sonnet 4.5 (Thinking)", + [Capability.TEXT, Capability.CODE, Capability.REASONING], + ), + ( + "claude-opus-4-5-thinking", + "Claude Opus 4.5 (Thinking)", + [Capability.TEXT, Capability.REASONING, Capability.CODE, Capability.LONG_CONTEXT], + ), + ("gpt-oss-120b-medium", "GPT-OSS 120B Medium", [Capability.TEXT, Capability.CODE]), +] + AGENT_TEMPLATES: list[tuple[str, RoleKey, str, list[str]]] = [ ( "مدیر سازمان", RoleKey.MANAGER, @@ -269,6 +308,92 @@ async def seed() -> None: ) logger.info("created mock provider + 2 models") + # The agy catalog (Google Antigravity). It is seeded for every + # deployment — the admin console needs the provider and its models to + # exist before any account is logged in. Routing activation is + # reconciled on *every* seed run below, not only at insert time, so + # flipping AI_DEFAULT_PROVIDER and re-seeding reshapes routing without + # a fresh database. + agy_is_default = settings.ai_default_provider.strip().lower() == "agy" + agy_provider_row = ( + await session.execute( + select(AiProviderRow).where(AiProviderRow.name == "agy") + ) + ).scalars().first() + if agy_provider_row is None: + agy_provider = Provider(name="agy", type=ProviderType.AGY, base_url="") + agy_provider_row = AiProviderRow( + id=agy_provider.id, name=agy_provider.name, + type=agy_provider.type.value, base_url="", is_active=True, + ) + session.add(agy_provider_row) + await session.flush() + for internal, display, caps in AGY_MODELS: + existing = ( + await session.execute( + select(AiModelRow).where( + AiModelRow.provider_id == agy_provider_row.id, + AiModelRow.internal_name == internal, + ) + ) + ).scalars().first() + if existing is not None: + continue + model = AiModel( + provider_id=agy_provider_row.id, + internal_name=internal, + display_name_for_admin=display, + capabilities=[str(c) for c in caps], + # agy is billed through the Google AI subscription, so the + # internal token cost estimate is zero; the operator sets the + # customer-facing price on the rate card. + input_price_per_1m=0.0, + output_price_per_1m=0.0, + ) + session.add( + AiModelRow( + id=model.id, provider_id=model.provider_id, + internal_name=model.internal_name, + display_name_for_admin=model.display_name_for_admin, + capabilities=model.capabilities, + input_price_per_1m=model.input_price_per_1m, + output_price_per_1m=model.output_price_per_1m, + is_active=False, + ) + ) + # Reconcile routing with the configured default on every run: agy models + # are routable only when agy is the default, and the mock catalog must + # not win when it is not — a model without accounts fails at call time. + agy_model_rows = ( + await session.execute( + select(AiModelRow).where( + AiModelRow.provider_id == agy_provider_row.id + ) + ) + ).scalars() + for row in agy_model_rows: + row.is_active = agy_is_default + if agy_is_default: + mock_provider_row = ( + await session.execute( + select(AiProviderRow).where(AiProviderRow.name == "mock") + ) + ).scalars().first() + if mock_provider_row is not None: + mock_model_rows = ( + await session.execute( + select(AiModelRow).where( + AiModelRow.provider_id == mock_provider_row.id + ) + ) + ).scalars() + for row in mock_model_rows: + row.is_active = False + logger.info( + "agy is the default provider; agy catalog activated " + "(accounts are added from the admin console)" + ) + if not (await session.execute(select(AgentTemplateRow))).scalars().first(): for name, role_key, description, tools in AGENT_TEMPLATES: template = AgentTemplate( diff --git a/apps/api/scripts/agy_login.py b/apps/api/scripts/agy_login.py new file mode 100644 index 0000000..f91c90d --- /dev/null +++ b/apps/api/scripts/agy_login.py @@ -0,0 +1,196 @@ +"""Add an agy (Google Antigravity) account through headless OAuth. + +The admin console does this over HTTP; this script is the same flow for an +operator sitting at the machine that runs the API — useful because the login +needs a browser redirect to a loopback URL that deliberately never loads. + + # 1) start a login; prints the Google URL to open + .venv/Scripts/python scripts/agy_login.py start + + # 2) after signing in, paste the browser's full address-bar URL back + .venv/Scripts/python scripts/agy_login.py finish "http://localhost:51121/oauth-callback?code=...&state=..." + +Run ``scripts/agy_login.py status`` to list stored accounts. The pending +challenge lives in ``.agy-login-state.json`` next to this script (gitignored) +for the ten minutes it is valid — never commit it. +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import sys +from pathlib import Path + +# The script lives in /apps/api/scripts/; app/ is one level up. +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from app.core.config import settings # noqa: E402 +from app.core.crypto import encrypt # noqa: E402 +from app.core.database import SessionFactory # noqa: E402 +from app.domain.ai_control.models import ( # noqa: E402 + ModelCredential, + Provider, + ProviderType, +) +from app.infrastructure.adapters.agy import ( # noqa: E402 + CHALLENGE_TTL_SECONDS, + AgyAuthClient, + AgyLoginChallenge, + AgyOAuthError, +) +from app.infrastructure.repositories.ai_control import ( # noqa: E402 + ModelCredentialRepository, + ProviderRepository, +) + +STATE_PATH = Path(__file__).resolve().parents[1] / ".agy-login-state.json" +AGY_PROVIDER_NAME = "agy" + + +def _save_challenge(challenge: AgyLoginChallenge, auth_url: str) -> None: + STATE_PATH.write_text( + json.dumps( + { + "login_id": challenge.login_id, + "auth_url": auth_url, + "redirect_uri": challenge.redirect_uri, + "verifier": challenge.verifier, + "state": challenge.state, + }, + ensure_ascii=False, + indent=2, + ), + encoding="utf-8", + ) + + +def _load_challenge() -> tuple[AgyLoginChallenge, str]: + if not STATE_PATH.exists(): + raise SystemExit( + "هیچ لاگین در انتظاری نیست. اول «start» را اجرا کنید.\n" + "No pending login. Run `start` first." + ) + data = json.loads(STATE_PATH.read_text(encoding="utf-8")) + challenge = AgyLoginChallenge( + login_id=data["login_id"], + auth_url=data["auth_url"], + redirect_uri=data["redirect_uri"], + verifier=data["verifier"], + state=data["state"], + created_at=0.0, + ) + return challenge, data["auth_url"] + + +async def cmd_start(_: argparse.Namespace) -> None: + challenge = AgyAuthClient().begin() + _save_challenge(challenge, challenge.auth_url) + print(challenge.auth_url) + print() + print("این لینک را در مرورگر باز کنید و با اکانت Google وارد شوید.") + print("Open this URL in a browser and sign in with the Google account.") + print( + f"بعد از تأیید، آدرس کامل نوار آدرس مرورگر را کپی کنید " + f"(با {challenge.redirect_uri} شروع می‌شود) و با «finish » بدهید." + ) + print(f"این درخواست {CHALLENGE_TTL_SECONDS // 60} دقیقه معتبر است.") + + +async def cmd_finish(args: argparse.Namespace) -> None: + challenge, _auth_url = _load_challenge() + try: + result = await AgyAuthClient().exchange(challenge, args.callback_url) + except AgyOAuthError as exc: + raise SystemExit(f"لاگین ناموفق بود: {exc}") from exc + if not result.email: + raise SystemExit("Google نشانی ایمیل اکانت را برنگرداند؛ لاگین لغو شد.") + + async with SessionFactory() as session: + providers = ProviderRepository(session) + credentials = ModelCredentialRepository(session) + provider = await providers.get_by_name(AGY_PROVIDER_NAME) + if provider is None: + provider = await providers.create( + Provider(name=AGY_PROVIDER_NAME, type=ProviderType.AGY, base_url="") + ) + envelope = json.dumps( + { + "refresh_token": result.refresh_token, + "project_id": result.project_id, + "email": result.email, + } + ) + existing = await credentials.find_by_email(provider.id, result.email) + if existing is not None: + await credentials.replace_material( + existing.id, encrypted_api_key=encrypt(envelope), + project_id=result.project_id, + ) + action = "به‌روزرسانی شد" + else: + await credentials.create( + ModelCredential( + provider_id=provider.id, + name=result.email, + encrypted_api_key=encrypt(envelope), + email=result.email, + project_id=result.project_id, + ) + ) + action = "اضافه شد" + await session.commit() + + STATE_PATH.unlink(missing_ok=True) + print(f"اکانت {result.email} {action}.") + print(f"project_id: {result.project_id or '(auto)'}") + print("توکن refresh به‌صورت رمزنگاری‌شده ذخیره شد و هرگز از API برنمی‌گردد.") + print( + "برای استفادهٔ واقعی، AI_DEFAULT_PROVIDER=agy را تنظیم و " + "«db:seed» را دوباره اجرا کنید." + ) + + +async def cmd_status(_: argparse.Namespace) -> None: + async with SessionFactory() as session: + providers = ProviderRepository(session) + provider = await providers.get_by_name(AGY_PROVIDER_NAME) + if provider is None: + print("provider agy وجود ندارد. یک لاگین شروع کنید.") + return + accounts = await ModelCredentialRepository(session).list_for_provider(provider.id) + if not accounts: + print("هیچ اکانتی ثبت نشده است.") + return + for account in accounts: + state = account.status if account.is_active else "disabled" + print(f"- {account.email or account.name} [{state}] last_used={account.last_used_at}") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + sub = parser.add_subparsers(dest="command", required=True) + sub.add_parser("start", help="شروع لاگین و چاپ لینک گوگل") + finish = sub.add_parser("finish", help="تکمیل لاگین با URL بازگشت مرورگر") + finish.add_argument("callback_url", help="آدرس کامل نوار آدرس بعد از لاگین") + sub.add_parser("status", help="فهرست اکانت‌های ثبت‌شده") + args = parser.parse_args() + + if settings.is_production and args.command in {"start", "finish"}: + print( + "هشدار: AI_DEFAULT_PROVIDER در محیط production است. " + "ادامه می‌دهید چون لاگین فقط credential می‌سازد." + ) + + asyncio.run( + { + "start": cmd_start, + "finish": cmd_finish, + "status": cmd_status, + }[args.command](args) + ) + + +if __name__ == "__main__": + main() diff --git a/apps/api/tests/test_agy.py b/apps/api/tests/test_agy.py new file mode 100644 index 0000000..bc3465c --- /dev/null +++ b/apps/api/tests/test_agy.py @@ -0,0 +1,715 @@ +"""agy provider tests — OAuth paste-URL login, translation, pool rotation. + +Everything external is faked at the ``httpx.AsyncClient`` seam, so the suite +runs with no Google, no network and no accounts. What is *not* faked is the +protocol shape: the fake answers are the shapes Google answers with (token +JSON, userinfo JSON, ``v1internal:generateContent`` responses, 429s), so a +regression in our translation fails here rather than in production. +""" + +from __future__ import annotations + +import json +import time +from urllib.parse import parse_qs, urlparse + +import pytest +import pytest_asyncio + +import app.infrastructure.adapters.agy as agy +from app.application.ports import ChatMessage +from app.core.errors import ProviderUnavailableError +from app.domain.ai_control.models import AiModel, ModelCredential +from tests.conftest import auth_header, signup + +# --------------------------------------------------------------------------- +# Fake httpx seam +# --------------------------------------------------------------------------- + +GOOD_GENERATE: dict = { + "candidates": [ + {"content": {"parts": [{"text": "سلام"}]}, "finishReason": "STOP"} + ], + "usageMetadata": { + "promptTokenCount": 7, + "candidatesTokenCount": 3, + "thoughtsTokenCount": 0, + "totalTokenCount": 10, + }, + "model": "gemini-3-flash", +} + + +class _FakeResponse: + def __init__( + self, + *, + status_code: int = 200, + json_data: dict | None = None, + text: str = "", + headers: dict | None = None, + ) -> None: + self.status_code = status_code + self._json = json_data if json_data is not None else {} + self._text = text + self.headers = headers or {} + + @property + def is_success(self) -> bool: + return self.status_code < 400 + + def json(self) -> dict: + return self._json + + @property + def text(self) -> str: + return self._text + + +def _token_response(refresh: str = "refresh-1", access: str = "access-1") -> _FakeResponse: + return _FakeResponse( + json_data={"access_token": access, "refresh_token": refresh, "expires_in": 3600} + ) + + +def patch_google(monkeypatch, handler) -> list[dict]: + """Point ``httpx.AsyncClient`` at a scripted handler. + + The adapter creates its own ``AsyncClient`` internally, so the seam is the + class itself. ``handler(method, url, json_body, headers, data)`` returns a + ``_FakeResponse`` (sync or async); every call is recorded and returned for + assertions. + """ + import inspect + + calls: list[dict] = [] + + async def _dispatch(method, url, json_body, headers, data): + result = handler(method, url, json_body, headers or {}, data or {}) + if inspect.isawaitable(result): + result = await result + return result + + async def _post(self, url, *, json=None, headers=None, data=None, **_): + calls.append({"method": "POST", "url": url, "json": json, "headers": headers, "data": data}) + return await _dispatch("POST", url, json, headers, data) + + async def _get(self, url, *, headers=None, **_): + calls.append({"method": "GET", "url": url, "json": None, "headers": headers, "data": None}) + return await _dispatch("GET", url, None, headers, None) + + class _Client: + def __init__(self, *_, **__): + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, *_) -> bool: + return False + + post = _post + get = _get + + monkeypatch.setattr(agy.httpx, "AsyncClient", _Client) + return calls + + +def default_handler(method, url, json_body, headers, data): + """The Google every happy-path test talks to.""" + if url.endswith("/token"): + return _token_response() + if "userinfo" in url: + return _FakeResponse(json_data={"email": "person@gmail.com"}) + if url.endswith("/v1internal:loadCodeAssist"): + return _FakeResponse(json_data={"cloudaicompanionProject": "agy-proj-123"}) + if url.endswith("/v1internal:fetchAvailableModels"): + return _FakeResponse(json_data={"models": {"gemini-3-flash": {}, "claude-sonnet-4-5": {}}}) + if url.endswith("/v1internal:generateContent"): + return _FakeResponse(json_data=GOOD_GENERATE) + raise AssertionError(f"unexpected {method} {url}") + + +@pytest.fixture +def google(monkeypatch): + """Async handler matching ``default_handler``; returns the call log.""" + return patch_google(monkeypatch, lambda *a: default_handler(*a)) + + +@pytest.fixture +def decrypted(monkeypatch): + """Make ``decrypt`` a passthrough so tests can build fake ciphertext.""" + monkeypatch.setattr(agy, "decrypt", lambda blob: blob.split("PLAINTEXT::", 1)[1]) + + +def _agy_credential( + *, email: str, refresh_token: str = "refresh-1", project_id: str = "agy-proj-123" +) -> ModelCredential: + """A credential as the adapter sees it: ciphertext already holds the blob.""" + envelope = json.dumps( + {"refresh_token": refresh_token, "project_id": project_id, "email": email} + ) + return ModelCredential( + provider_id="prov-1", + name=email, + encrypted_api_key=f"PLAINTEXT::{envelope}", + email=email, + project_id=project_id, + ) + + +_MODEL = AiModel(provider_id="p", internal_name="gemini-3-flash", display_name_for_admin="f") +_MESSAGES = [ChatMessage(role="user", content="سلام")] + + +# --------------------------------------------------------------------------- +# HTTP endpoints (platform owner only) +# --------------------------------------------------------------------------- + + +async def test_agy_endpoints_require_platform_owner(client, google): + body = await signup(client, "customer@example.com") + response = await client.post( + "/api/admin/agy/accounts/login/start", headers=auth_header(body["access_token"]) + ) + assert response.status_code == 403 + + +async def test_agy_login_flow_end_to_end(client, google, platform_owner_token): + # 1. start the headless login. + started = await client.post( + "/api/admin/agy/accounts/login/start", + headers=auth_header(platform_owner_token), + ) + assert started.status_code == 201, started.text + login_id = started.json()["login_id"] + auth_url = started.json()["auth_url"] + state = parse_qs(urlparse(auth_url).query)["state"][0] + + # 2. the operator pastes the browser's callback URL back. + completed = await client.post( + "/api/admin/agy/accounts/login/complete", + headers=auth_header(platform_owner_token), + json={ + "login_id": login_id, + "callback_url": ( + f"http://localhost:51121/oauth-callback?code=auth-code&state={state}" + ), + }, + ) + assert completed.status_code == 201, completed.text + account = completed.json() + assert account["email"] == "person@gmail.com" + assert account["status"] == "active" + assert account["is_active"] is True + + # No credential material ever crosses the API surface. + assert "refresh_token" not in completed.text + assert "encrypted" not in completed.text + assert "agy-proj-123" in completed.text # project id is admin metadata + + # 3. the account is listed. + listed = await client.get( + "/api/admin/agy/accounts", headers=auth_header(platform_owner_token) + ) + assert listed.status_code == 200 + emails = [a["email"] for a in listed.json()] + assert emails == ["person@gmail.com"] + + +async def test_agy_complete_login_rejects_unrelated_url(client, google, platform_owner_token): + started = await client.post( + "/api/admin/agy/accounts/login/start", + headers=auth_header(platform_owner_token), + ) + completed = await client.post( + "/api/admin/agy/accounts/login/complete", + headers=auth_header(platform_owner_token), + json={ + "login_id": started.json()["login_id"], + "callback_url": "http://localhost:51121/oauth-callback?code=x&state=elsewhere", + }, + ) + assert completed.status_code == 422 + # The failed attempt must not have created an account. + listed = await client.get( + "/api/admin/agy/accounts", headers=auth_header(platform_owner_token) + ) + assert listed.json() == [] + + +async def test_agy_account_test_and_delete(client, google, platform_owner_token): + started = await client.post( + "/api/admin/agy/accounts/login/start", + headers=auth_header(platform_owner_token), + ) + state = parse_qs(urlparse(started.json()["auth_url"]).query)["state"][0] + created = await client.post( + "/api/admin/agy/accounts/login/complete", + headers=auth_header(platform_owner_token), + json={ + "login_id": started.json()["login_id"], + "callback_url": ( + f"http://localhost:51121/oauth-callback?code=c&state={state}" + ), + }, + ) + account_id = created.json()["id"] + + # A test call runs a real (faked-Google) completion through the account. + tested = await client.post( + f"/api/admin/agy/accounts/{account_id}/test", + headers=auth_header(platform_owner_token), + ) + assert tested.status_code == 200, tested.text + assert tested.json()["ok"] is True + assert tested.json()["detail"] == "سلام" + + deleted = await client.delete( + f"/api/admin/agy/accounts/{account_id}", + headers=auth_header(platform_owner_token), + ) + assert deleted.status_code == 204 + listed = await client.get( + "/api/admin/agy/accounts", headers=auth_header(platform_owner_token) + ) + assert listed.json() == [] + + +async def test_agy_same_account_twice_replaces(client, google, platform_owner_token): + """Re-login of a known Google account rotates the key, never duplicates.""" + for _ in range(2): + started = await client.post( + "/api/admin/agy/accounts/login/start", + headers=auth_header(platform_owner_token), + ) + state = parse_qs(urlparse(started.json()["auth_url"]).query)["state"][0] + done = await client.post( + "/api/admin/agy/accounts/login/complete", + headers=auth_header(platform_owner_token), + json={ + "login_id": started.json()["login_id"], + "callback_url": ( + f"http://localhost:51121/oauth-callback?code=c&state={state}" + ), + }, + ) + assert done.status_code == 201 + listed = await client.get( + "/api/admin/agy/accounts", headers=auth_header(platform_owner_token) + ) + assert len(listed.json()) == 1 + + +# --------------------------------------------------------------------------- +# Seed: the agy catalog +# --------------------------------------------------------------------------- + + +@pytest_asyncio.fixture +async def seeded_install(monkeypatch, engine, session_factory): + """Run the real ``app.seed`` against the in-memory test database.""" + import app.seed as seed_module + + monkeypatch.setattr(seed_module, "engine", engine) + monkeypatch.setattr(seed_module, "SessionFactory", session_factory) + await seed_module.seed() + return seed_module + + +async def _agy_models(session_factory) -> list[tuple[str, bool]]: + from sqlalchemy import select + + from app.infrastructure.db.models import AiModelRow, AiProviderRow + + async with session_factory() as session: + rows = ( + await session.execute( + select(AiModelRow, AiProviderRow.name) + .join(AiProviderRow, AiModelRow.provider_id == AiProviderRow.id) + .where(AiProviderRow.name == "agy") + ) + ).all() + return [(row.internal_name, row.is_active) for row, _name in rows] + + +async def test_seed_creates_agy_catalog(client, seeded_install, session_factory): + models = await _agy_models(session_factory) + names = {name for name, _ in models} + assert names == { + "gemini-3-flash", + "gemini-3-pro-high", + "claude-sonnet-4-5", + "claude-sonnet-4-5-thinking", + "claude-opus-4-5-thinking", + "gpt-oss-120b-medium", + } + # Default deployment keeps routing on mock: agy models wait inactive. + assert all(active is False for _, active in models) + + +async def test_seed_activates_agy_when_it_is_default( + monkeypatch, client, seeded_install, session_factory +): + from app.core.config import settings + + monkeypatch.setattr(settings, "ai_default_provider", "agy") + await seeded_install.seed() + models = await _agy_models(session_factory) + assert all(active is True for _, active in models) + + +async def test_seed_deactivates_mock_when_agy_is_default( + monkeypatch, client, seeded_install, session_factory +): + from sqlalchemy import select + + from app.core.config import settings + from app.infrastructure.db.models import AiModelRow, AiProviderRow + + monkeypatch.setattr(settings, "ai_default_provider", "agy") + await seeded_install.seed() + + async with session_factory() as session: + rows = ( + await session.execute( + select(AiModelRow) + .join(AiProviderRow, AiModelRow.provider_id == AiProviderRow.id) + .where(AiProviderRow.name == "mock") + ) + ).scalars() + assert all(row.is_active is False for row in rows) + + +# --------------------------------------------------------------------------- +# OAuth login flow +# --------------------------------------------------------------------------- + + +def test_auth_url_shape(): + challenge = agy.AgyAuthClient().begin() + query = parse_qs(urlparse(challenge.auth_url).query) + assert challenge.auth_url.startswith( + "https://accounts.google.com/o/oauth2/v2/auth" + ) + assert query["client_id"][0] == ( + "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com" + ) + assert query["redirect_uri"][0] == "http://localhost:51121/oauth-callback" + assert query["code_challenge_method"][0] == "S256" + # access_type=offline + prompt=consent is what makes Google hand back a + # refresh token; without it the login completes but nothing can be stored. + assert query["access_type"][0] == "offline" + assert query["prompt"][0] == "consent" + assert "cloud-platform" in query["scope"][0] + assert challenge.verifier and challenge.state + + +async def test_paste_flow_exchanges_and_resolves_account(google): + """start → paste callback URL → account material, verifier stays server-side.""" + login_id, _auth_url = await agy.begin_login() + challenge = agy._CHALLENGES[login_id] + callback = ( + f"http://localhost:51121/oauth-callback?code=auth-code&state={challenge.state}" + ) + result = await agy.redeem_login(login_id, callback) + assert result.email == "person@gmail.com" + assert result.project_id == "agy-proj-123" + assert result.refresh_token == "refresh-1" + # The challenge is single-use. + with pytest.raises(agy.AgyOAuthError): + await agy.redeem_login(login_id, callback) + + +async def test_foreign_state_is_rejected(google): + login_id, _ = await agy.begin_login() + with pytest.raises(agy.AgyOAuthError): + await agy.redeem_login( + login_id, "http://localhost:51121/oauth-callback?code=x&state=not-ours" + ) + + +async def test_missing_code_is_rejected(google): + login_id, _ = await agy.begin_login() + with pytest.raises(agy.AgyOAuthError): + await agy.redeem_login(login_id, "http://localhost:51121/oauth-callback?state=x") + + +async def test_google_error_page_is_rejected(google): + login_id, _ = await agy.begin_login() + with pytest.raises(agy.AgyOAuthError): + await agy.redeem_login( + login_id, "http://localhost:51121/oauth-callback?error=access_denied" + ) + + +async def test_refresh_token_missing_is_actionable(monkeypatch): + """A token response without ``refresh_token`` fails loudly and tells why.""" + async def handler(method, url, json_body, headers, data): + if url.endswith("/token"): + return _FakeResponse( + json_data={"access_token": "a", "expires_in": 3600} + ) + return await default_handler(method, url, json_body, headers, data) + + patch_google(monkeypatch, handler) + login_id, _ = await agy.begin_login() + challenge = agy._CHALLENGES[login_id] + with pytest.raises(agy.AgyOAuthError) as excinfo: + await agy.redeem_login( + login_id, + f"http://localhost:51121/oauth-callback?code=c&state={challenge.state}", + ) + assert "refresh_token" in str(excinfo.value) + + +async def test_rejected_refresh_is_an_account_error_not_an_outage(monkeypatch): + async def handler(method, url, json_body, headers, data): + if url.endswith("/token"): + return _FakeResponse( + status_code=400, + text=json.dumps({"error": "invalid_grant"}), + ) + return await default_handler(method, url, json_body, headers, data) + + patch_google(monkeypatch, handler) + with pytest.raises(Exception) as excinfo: + await agy.AgyAuthClient().refresh("dead-token") + # Raised as the internal taxonomy so the adapter benches the account. + assert type(excinfo.value).__name__ == "_AgyAuthError" + + +# --------------------------------------------------------------------------- +# Request translation (pure functions) +# --------------------------------------------------------------------------- + + +def test_payload_translation(decrypted): + account = agy.AgyAccount( + credential_id="c1", email="a@b.c", project_id="proj-1", refresh_token="refresh-1" + ) + payload = agy._build_generate_payload( + account, + "gemini-3-pro-high", + [ + ChatMessage(role="system", content="تو مدیر هستی."), + ChatMessage(role="user", content="سلام"), + ChatMessage(role="assistant", content="سلام! چیزی کمکی می‌خواهید؟"), + ChatMessage(role="user", content="یک گزارش بساز"), + ], + system="پرامپت سیستم", + temperature=0.4, + max_tokens=512, + ) + assert payload["model"] == "gemini-3-pro-high" + assert payload["project"] == "proj-1" + assert payload["requestId"].startswith("agent-") + assert payload["request"]["sessionId"].startswith("-") + # System content from both the `system` argument and system messages lands + # in systemInstruction, not as a user turn. + instruction = payload["request"]["systemInstruction"]["parts"][0]["text"] + assert "پرامپت سیستم" in instruction + assert "تو مدیر هستی." in instruction + roles = [c["role"] for c in payload["request"]["contents"]] + assert roles == ["user", "model", "user"] + assert payload["request"]["generationConfig"]["temperature"] == 0.4 + assert payload["request"]["generationConfig"]["maxOutputTokens"] == 512 + # Function calling must be VALIDATED or tool schemas get rejected. + assert payload["request"]["toolConfig"]["functionCallingConfig"]["mode"] == "VALIDATED" + assert all(s["threshold"] == "OFF" for s in payload["request"]["safetySettings"]) + + +def test_response_parsing_skips_thoughts(): + model = AiModel(provider_id="p", internal_name="m", display_name_for_admin="m") + response = agy._parse_generate_response( + { + "candidates": [ + { + "content": { + "parts": [ + {"text": "دارم فکر می‌کنم", "thought": True}, + {"text": "پاسخ نهایی"}, + ] + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 11, + "candidatesTokenCount": 5, + "thoughtsTokenCount": 4, + }, + }, + model=model, + started=0.0, + ) + assert response.content == "پاسخ نهایی" + assert response.input_tokens == 11 + assert response.output_tokens == 5 + assert response.raw["thought_tokens"] == 4 + assert response.raw["finish_reason"] == "STOP" + + +# --------------------------------------------------------------------------- +# The adapter: completion + account pool rotation +# --------------------------------------------------------------------------- + + +async def test_complete_happy_path(google, decrypted): + adapter = agy.AgyModelAdapter() + response = await adapter.complete( + model=_MODEL, credentials=[_agy_credential(email="one@gmail.com")], + messages=_MESSAGES, + ) + assert response.content == "سلام" + assert response.input_tokens == 7 + assert response.output_tokens == 3 + generate = [c for c in google if c["url"].endswith("generateContent")] + assert len(generate) == 1 + # The account's project id is what bills the call. + assert generate[0]["json"]["project"] == "agy-proj-123" + + +async def test_rate_limited_account_is_rotated_and_benched(monkeypatch, decrypted): + """A 429 on one account must not fail the request: the pool moves on.""" + async def handler(method, url, json_body, headers, data): + if url.endswith("/token"): + # Each account's refresh token maps to its own access token. + access = {"refresh-1": "access-one", "refresh-2": "access-two"}[ + data.get("refresh_token", "") + ] + return _token_response(refresh=data.get("refresh_token", ""), access=access) + if url.endswith("/v1internal:generateContent"): + auth = headers.get("Authorization", "") + if auth.endswith("access-one"): + # Both base URLs hit the same quota. + return _FakeResponse( + status_code=429, text="quota", headers={"Retry-After": "120"} + ) + return _FakeResponse(json_data=GOOD_GENERATE) + return await default_handler(method, url, json_body, headers, data) + + patch_google(monkeypatch, handler) + reported: list[tuple] = [] + + async def report(credential_id, *, status=None, cooldown_until=None, error=None, used=False): + reported.append((credential_id, status, cooldown_until, error, used)) + + adapter = agy.AgyModelAdapter(report=report) + # Start the rotation cursor so `one` is genuinely attempted first. + adapter._state.cursor = -1 + one = _agy_credential(email="one@gmail.com", refresh_token="refresh-1") + two = _agy_credential(email="two@gmail.com", refresh_token="refresh-2") + response = await adapter.complete( + model=_MODEL, credentials=[one, two], messages=_MESSAGES + ) + assert response.content == "سلام" + + cooldowns = [r for r in reported if r[1] == "cooldown"] + assert cooldowns, "the rate-limited account must be reported" + assert cooldowns[0][0] == one.id + # Google's Retry-After is honoured: the bench is ~120s out, not the + # configured fallback. + remaining = cooldowns[0][2].timestamp() - time.time() + assert 90 <= remaining <= 200 + + +async def test_every_account_rate_limited_raises(monkeypatch, decrypted): + async def handler(method, url, json_body, headers, data): + if url.endswith("/token"): + return _token_response() + if url.endswith("/v1internal:generateContent"): + return _FakeResponse(status_code=429, text="quota") + return await default_handler(method, url, json_body, headers, data) + + patch_google(monkeypatch, handler) + adapter = agy.AgyModelAdapter() + with pytest.raises(ProviderUnavailableError): + await adapter.complete( + model=_MODEL, + credentials=[ + _agy_credential(email="one@gmail.com", refresh_token="r1"), + _agy_credential(email="two@gmail.com", refresh_token="r2"), + ], + messages=_MESSAGES, + ) + + +async def test_rejected_grant_marks_account_error(monkeypatch, decrypted): + """401 means the OAuth grant died — bench as error, move to next account.""" + async def handler(method, url, json_body, headers, data): + if url.endswith("/token"): + access = {"refresh-1": "access-one", "refresh-2": "access-two"}[ + data.get("refresh_token", "") + ] + return _token_response(refresh=data.get("refresh_token", ""), access=access) + if url.endswith("/v1internal:generateContent"): + auth = headers.get("Authorization", "") + if auth.endswith("access-one"): + return _FakeResponse(status_code=401, text="unauthorized") + return _FakeResponse(json_data=GOOD_GENERATE) + return await default_handler(method, url, json_body, headers, data) + + patch_google(monkeypatch, handler) + reported: list[tuple] = [] + + async def report(credential_id, *, status=None, cooldown_until=None, error=None, used=False): + reported.append((credential_id, status)) + + adapter = agy.AgyModelAdapter(report=report) + adapter._state.cursor = -1 + dead = _agy_credential(email="dead@gmail.com", refresh_token="refresh-1") + alive = _agy_credential(email="alive@gmail.com", refresh_token="refresh-2") + response = await adapter.complete( + model=_MODEL, credentials=[dead, alive], messages=_MESSAGES + ) + assert response.content == "سلام" + errors = [r for r in reported if r[1] == "error"] + assert errors and errors[0][0] == dead.id + # The healthy account is cleared from any bench and marked used. + assert any(r[0] == alive.id and r[1] == "active" for r in reported) + + +async def test_server_error_falls_through_to_next_backend(monkeypatch, decrypted): + """5xx on one backend URL is not the account's fault — try the other.""" + async def handler(method, url, json_body, headers, data): + if url.endswith("/token"): + return _token_response() + if url.endswith("/v1internal:generateContent"): + if "daily-cloudcode-pa" in url: + return _FakeResponse(status_code=503, text="backend down") + return _FakeResponse(json_data=GOOD_GENERATE) + return await default_handler(method, url, json_body, headers, data) + + patch_google(monkeypatch, handler) + adapter = agy.AgyModelAdapter() + response = await adapter.complete( + model=_MODEL, credentials=[_agy_credential(email="one@gmail.com")], + messages=_MESSAGES, + ) + assert response.content == "سلام" + + +async def test_no_accounts_raises_actionable_error(google, decrypted): + adapter = agy.AgyModelAdapter() + with pytest.raises(ProviderUnavailableError) as excinfo: + await adapter.complete(model=_MODEL, credentials=[], messages=_MESSAGES) + assert "اکانت agy" in str(excinfo.value) + + +async def test_health_reports_account_count(google, decrypted): + async def loader(): + return [_agy_credential(email="one@gmail.com")] + + adapter = agy.AgyModelAdapter(loader=loader) + health = await adapter.health() + assert health.ok and "1" in health.detail + + +async def test_available_models(google, decrypted): + async def loader(): + return [_agy_credential(email="one@gmail.com")] + + adapter = agy.AgyModelAdapter(loader=loader) + models = await adapter.available_models() + assert models == ["claude-sonnet-4-5", "gemini-3-flash"] diff --git a/apps/web/src/app/(app)/admin/agy-accounts/page.tsx b/apps/web/src/app/(app)/admin/agy-accounts/page.tsx new file mode 100644 index 0000000..8638302 --- /dev/null +++ b/apps/web/src/app/(app)/admin/agy-accounts/page.tsx @@ -0,0 +1,393 @@ +"use client"; + +import * as React from "react"; +import type { AgyAccount } from "@agy/contracts"; +import { + Badge, + Button, + Card, + CardHint, + CardTitle, + ConfirmDialog, + EmptyState, + ErrorState, + Field, + LoadingState, + Modal, + ToastStack, + type ToastMessage, +} from "@agy/ui"; +import { ApiError, api } from "@/lib/api"; + +/** + * agy account management (platform owner only). + * + * agy has no API keys — an account is a Google identity. The login flow is the + * headless paste-URL shape of Google OAuth: + * + * 1. start login → the API returns an auth URL; + * 2. the owner opens it, signs in, and Google redirects to a loopback URL + * that deliberately never finishes loading — the code and state are + * already in the address bar; + * 3. the owner pastes that whole URL back and the API redeems it. + * + * The rotation pool spreads calls across the accounts listed here and benches + * a rate-limited one until its quota window resets, so adding more accounts is + * the sanctioned way to raise agy throughput. + */ + +const STATUS_LABEL: Record = { + active: "فعال", + cooldown: "استراحت (سهمیه)", + error: "خطا — نیاز به لاگین مجدد", +}; + +export default function AdminAgyAccountsPage() { + const [accounts, setAccounts] = React.useState([]); + const [loading, setLoading] = React.useState(true); + const [error, setError] = React.useState(null); + const [loginOpen, setLoginOpen] = React.useState(false); + const [toDelete, setToDelete] = React.useState(null); + const [toasts, setToasts] = React.useState([]); + + const pushToast = React.useCallback( + (title: string, description?: string) => { + const id = Math.random().toString(36).slice(2); + setToasts((current) => [...current, { id, title, description }]); + window.setTimeout( + () => setToasts((current) => current.filter((t) => t.id !== id)), + 6000, + ); + }, + [], + ); + + const load = React.useCallback(async () => { + setLoading(true); + setError(null); + try { + setAccounts(await api.agyAccounts()); + } catch (err) { + setError(err instanceof ApiError ? err.message : "دریافت اکانت‌ها ناموفق بود."); + } finally { + setLoading(false); + } + }, []); + + React.useEffect(() => { + void load(); + }, [load]); + + const onAccountAdded = (account: AgyAccount) => { + setAccounts((current) => { + const exists = current.some((a) => a.id === account.id); + return exists + ? current.map((a) => (a.id === account.id ? account : a)) + : [account, ...current]; + }); + setLoginOpen(false); + pushToast("اکانت agy اضافه شد", account.email || undefined); + }; + + const onSetActive = async (account: AgyAccount, active: boolean) => { + try { + const updated = await api.agySetAccountActive(account.id, active); + setAccounts((current) => current.map((a) => (a.id === updated.id ? updated : a))); + } catch (err) { + pushToast( + "تغییر وضعیت ناموفق", + err instanceof ApiError ? err.message : undefined, + ); + } + }; + + const onDelete = async () => { + if (!toDelete) return; + try { + await api.agyDeleteAccount(toDelete.id); + setAccounts((current) => current.filter((a) => a.id !== toDelete.id)); + pushToast("اکانت حذف شد", toDelete.email || undefined); + } catch (err) { + pushToast("حذف ناموفق", err instanceof ApiError ? err.message : undefined); + } finally { + setToDelete(null); + } + }; + + return ( +
+
+
+

اکانت‌های agy

+

+ agy کلید API ندارد: هر اکانت یک هویت Google است که از طریق OAuth وارد + می‌شود. هر تعداد اکانت اضافه کنید؛ پول چرخشی بار را بین آن‌ها پخش + می‌کند و اکانت محدودشده را تا پایان پنجرهٔ سهمیه کنار می‌گذارد. +

+
+ +
+ + {loading ? ( + + ) : error ? ( + + ) : accounts.length === 0 ? ( + + ) : ( +
+ {accounts.map((account) => ( + setToDelete(account)} + /> + ))} +
+ )} + + {loginOpen && ( + setLoginOpen(false)} + onDone={onAccountAdded} + onError={(message) => pushToast("لاگین ناموفق", message)} + /> + )} + + setToDelete(null)} + /> + + + setToasts((current) => current.filter((t) => t.id !== id)) + } /> +
+ ); +} + +function AgyAccountCard({ + account, + onSetActive, + onDelete, +}: { + account: AgyAccount; + onSetActive: (account: AgyAccount, active: boolean) => void; + onDelete: () => void; +}) { + const [testing, setTesting] = React.useState(false); + const [testResult, setTestResult] = React.useState(null); + + const onTest = async () => { + setTesting(true); + setTestResult(null); + try { + const result = await api.agyTestAccount(account.id); + setTestResult( + result.ok ? `پاسخ دریافت شد: ${result.detail}` : `تست ناموفق: ${result.detail}`, + ); + } catch (err) { + setTestResult(err instanceof ApiError ? err.message : "تست ناموفق بود."); + } finally { + setTesting(false); + } + }; + + const cooldownRemaining = React.useMemo(() => { + if (account.status !== "cooldown" || !account.cooldown_until) return null; + const ms = new Date(account.cooldown_until).getTime() - Date.now(); + if (ms <= 0) return null; + return Math.ceil(ms / 60000); + }, [account.status, account.cooldown_until]); + + return ( + +
+
+
+

+ {account.email || account.name} +

+ + {account.is_active ? "فعال" : "غیرفعال"} + + + {STATUS_LABEL[account.status]} + {cooldownRemaining ? ` · ${cooldownRemaining} دقیقه` : ""} + +
+

+ {account.last_used_at + ? `آخرین استفاده: ${new Date(account.last_used_at).toLocaleString("fa-IR")}` + : "هنوز استفاده نشده"} + {account.last_error ? ` · ${account.last_error}` : ""} +

+ {testResult &&

{testResult}

} +
+
+ + + +
+
+ + ); +} + +function LoginModal({ + onClose, + onDone, + onError, +}: { + onClose: () => void; + onDone: (account: AgyAccount) => void; + onError: (message: string) => void; +}) { + const [step, setStep] = React.useState<"start" | "redeem">("start"); + const [authUrl, setAuthUrl] = React.useState(""); + const [loginId, setLoginId] = React.useState(""); + const [callbackUrl, setCallbackUrl] = React.useState(""); + const [busy, setBusy] = React.useState(false); + const [formError, setFormError] = React.useState(null); + + const onStart = async () => { + setBusy(true); + setFormError(null); + try { + const started = await api.agyStartLogin(); + setLoginId(started.login_id); + setAuthUrl(started.auth_url); + setStep("redeem"); + } catch (err) { + setFormError(err instanceof ApiError ? err.message : "شروع لاگین ناموفق بود."); + } finally { + setBusy(false); + } + }; + + const onComplete = async (event: React.FormEvent) => { + event.preventDefault(); + setBusy(true); + setFormError(null); + try { + const account = await api.agyCompleteLogin({ + login_id: loginId, + callback_url: callbackUrl.trim(), + }); + onDone(account); + } catch (err) { + const message = err instanceof ApiError ? err.message : "تکمیل لاگین ناموفق بود."; + setFormError(message); + onError(message); + } finally { + setBusy(false); + } + }; + + return ( + + + + + ) : ( + + ) + } + > + {step === "start" ? ( +
+

+ ۱. روی «دریافت لینک لاگین» بزنید. یک آدرس احراز هویت گوگل ساخته + می‌شود که فقط برای این درخواست معتبر است (۱۰ دقیقه). +

+

+ ۲. آن آدرس را در مرورگر باز کنید و با اکانت Google مورد نظر وارد + شوید. +

+

+ ۳. بعد از تأیید، مرورگر به یک آدرس محلی منتقل می‌شود که عمداً + بارگذاری نمی‌شود. نوار آدرس را همان‌جا کپی کنید (آدرس کامل با + ?code=…&state=…) + و در گام بعدی Paste کنید. +

+
+ ) : ( +
+
+

+ این آدرس را در مرورگر باز کنید و با Google وارد شوید: +

+ + {authUrl} + +

+ بعد از تأیید در گوگل، آدرس کامل نوار آدرس مرورگر (که با + oauth-callback?… شروع + می‌شود) را در کادر زیر بچسبانید. +

+
+
+ setCallbackUrl(e.target.value)} + placeholder="http://localhost:51121/oauth-callback?code=…&state=…" + className="ltr-inline" + autoComplete="off" + /> + {formError && } + +
+ )} + + ); +} diff --git a/apps/web/src/app/(app)/admin/layout.tsx b/apps/web/src/app/(app)/admin/layout.tsx index 9809707..e146473 100644 --- a/apps/web/src/app/(app)/admin/layout.tsx +++ b/apps/web/src/app/(app)/admin/layout.tsx @@ -16,6 +16,7 @@ import { useSession } from "@/lib/session"; const SECTIONS = [ { href: "/admin", label: "نمای کلی" }, + { href: "/admin/agy-accounts", label: "اکانت‌های agy" }, { href: "/admin/providers", label: "ارائه‌دهنده‌ها" }, { href: "/admin/models", label: "مدل‌ها" }, { href: "/admin/routing", label: "مسیریابی" }, diff --git a/apps/web/src/app/(app)/admin/providers/page.tsx b/apps/web/src/app/(app)/admin/providers/page.tsx index 6ee98a2..3917a51 100644 --- a/apps/web/src/app/(app)/admin/providers/page.tsx +++ b/apps/web/src/app/(app)/admin/providers/page.tsx @@ -22,6 +22,7 @@ const PROVIDER_TYPES: { value: ProviderType; label: string }[] = [ { value: "anthropic", label: "Anthropic" }, { value: "gemini", label: "Gemini" }, { value: "openrouter", label: "OpenRouter" }, + { value: "agy", label: "agy (Google Antigravity)" }, ]; export default function AdminProvidersPage() { diff --git a/apps/web/src/lib/api.ts b/apps/web/src/lib/api.ts index da5e985..f6ae2f5 100644 --- a/apps/web/src/lib/api.ts +++ b/apps/web/src/lib/api.ts @@ -438,6 +438,36 @@ export const api = { `/api/admin/providers/${providerId}/credentials`, { method: "POST", body }, ), + // agy accounts: OAuth login instead of API keys. + agyAccounts: () => + apiRequest("/api/admin/agy/accounts"), + agyStartLogin: () => + apiRequest<{ login_id: string; auth_url: string }>( + "/api/admin/agy/accounts/login/start", + { method: "POST" }, + ), + agyCompleteLogin: (body: { login_id: string; callback_url: string }) => + apiRequest( + "/api/admin/agy/accounts/login/complete", + { method: "POST", body }, + ), + agySetAccountActive: (accountId: string, active: boolean) => + apiRequest( + `/api/admin/agy/accounts/${accountId}/${active ? "activate" : "deactivate"}`, + { method: "POST" }, + ), + agyDeleteAccount: (accountId: string) => + apiRequest(`/api/admin/agy/accounts/${accountId}`, { method: "DELETE" }), + agyTestAccount: (accountId: string) => + apiRequest<{ + ok: boolean; + detail: string; + model_internal_name: string; + input_tokens: number; + output_tokens: number; + }>(`/api/admin/agy/accounts/${accountId}/test`, { method: "POST" }), + agyModels: () => + apiRequest("/api/admin/agy/models"), adminModels: () => apiRequest("/api/admin/models"), adminCreateModel: (body: { provider_id: string; diff --git a/apps/web/src/test/api.test.ts b/apps/web/src/test/api.test.ts index 7070524..b627c92 100644 --- a/apps/web/src/test/api.test.ts +++ b/apps/web/src/test/api.test.ts @@ -111,3 +111,60 @@ describe("error console client", () => { expect(calls[0]).not.toContain("?"); }); }); + +describe("agy account client", () => { + afterEach(() => { + vi.unstubAllGlobals(); + }); + + function stubJson(body: unknown) { + const calls: { method: string; url: string; body: string }[] = []; + vi.stubGlobal( + "fetch", + vi.fn(async (url: string, init?: RequestInit) => { + calls.push({ + method: init?.method ?? "GET", + url, + body: String(init?.body ?? ""), + }); + return new Response(JSON.stringify(body), { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + }), + ); + return calls; + } + + it("posts to the login start endpoint", async () => { + const calls = stubJson({ login_id: "l1", auth_url: "https://accounts.google.com/..." }); + await api.agyStartLogin(); + expect(calls[0].method).toBe("POST"); + expect(calls[0].url).toContain("/api/admin/agy/accounts/login/start"); + }); + + it("sends the pasted callback URL back verbatim", async () => { + const calls = stubJson({ id: "a1", email: "x@gmail.com" }); + await api.agyCompleteLogin({ + login_id: "l1", + callback_url: "http://localhost:51121/oauth-callback?code=c&state=s", + }); + expect(calls[0].url).toContain("/api/admin/agy/accounts/login/complete"); + expect(calls[0].body).toContain("callback_url"); + }); + + it("routes activate/deactivate to the right verb", async () => { + const calls = stubJson({ id: "a1" }); + await api.agySetAccountActive("a1", false); + await api.agySetAccountActive("a1", true); + expect(calls[0].url).toContain("/api/admin/agy/accounts/a1/deactivate"); + expect(calls[1].url).toContain("/api/admin/agy/accounts/a1/activate"); + }); + + it("deletes with the DELETE verb", async () => { + const calls = stubJson(null); + await api.agyDeleteAccount("a1"); + expect(calls[0].method).toBe("DELETE"); + expect(calls[0].url).toContain("/api/admin/agy/accounts/a1"); + }); +}); diff --git a/packages/contracts/src/index.ts b/packages/contracts/src/index.ts index eca3804..3ccc467 100644 --- a/packages/contracts/src/index.ts +++ b/packages/contracts/src/index.ts @@ -16,7 +16,29 @@ export type ProviderType = | "openai_compatible" | "anthropic" | "gemini" - | "openrouter"; + | "openrouter" + | "agy"; + +/** Rotation-pool lifecycle of one agy account. */ +export type AgyAccountStatus = "active" | "cooldown" | "error"; + +/** + * One agy account (Google identity). Deliberately credential-free: the console + * manages accounts, never tokens. + */ +export interface AgyAccount { + id: string; + provider_id: string; + name: string; + email: string; + project_id: string; + status: AgyAccountStatus; + cooldown_until: string | null; + is_active: boolean; + last_error: string; + last_used_at: string | null; + created_at: string; +} /** What a model can do. Routing picks models by required capability. */ export type Capability = "text" | "vision" | "tools" | "reasoning" | "long_context" | "code";