380 lines
13 KiB
Python
380 lines
13 KiB
Python
"""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
|
|
|
|
|
|
class AgyAccountModelOut(BaseModel):
|
|
id: str
|
|
display_name: str
|
|
supports_thinking: bool = False
|
|
supported_efforts: list[str] = Field(default_factory=list)
|
|
default_effort: str | None = None
|
|
supports_images: bool = False
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 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("/accounts/{credential_id}/models", response_model=list[AgyAccountModelOut])
|
|
async def list_account_models(
|
|
credential_id: str,
|
|
session: AsyncSession = Depends(get_session),
|
|
) -> list[AgyAccountModelOut]:
|
|
"""Live model catalog of one account — fetched from Google on every call.
|
|
|
|
The admin console calls this when an account is selected, so the model
|
|
list always reflects that subscription's current entitlement along with
|
|
the supported reasoning efforts for each model.
|
|
"""
|
|
from app.core.errors import ValidationError
|
|
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: one forced account
|
|
try:
|
|
models = await adapter.available_models_for(credential)
|
|
return [
|
|
AgyAccountModelOut(
|
|
id=m.id,
|
|
display_name=m.display_name,
|
|
supports_thinking=m.supports_thinking,
|
|
supported_efforts=m.supported_efforts,
|
|
default_effort=m.default_effort,
|
|
supports_images=m.supports_images,
|
|
)
|
|
for m in models
|
|
]
|
|
except ValidationError:
|
|
raise
|
|
except Exception as exc: # noqa: BLE001 - the console wants the reason
|
|
raise ValidationError(str(exc)) from exc
|