u78649634_default/smart-org
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

افزودن ارائه‌دهندهٔ agy (Google Antigravity) با لاگین headless و پول چرخشی اکانت

Browse Source 
- Adapter اکانت‌محور با چرخش round-robin و بنچ‌شدن خودکار روی ۴۲۹/۴۰۱
- ورود ادمین از طریق URL بازگشتی گوگل (PKCE) و ذخیرهٔ رمزنگاری‌شدهٔ refresh token
- صفحهٔ مدیریت اکانت‌های agy در پنل ادمین به همراه تست و فعال/غیرفعال‌سازی
- مهاجرت Alembic b1c2d3e4f5a6 برای ستون‌های اکانت agy
- seed همیشه کاتالوگ agy را می‌سازد و فقط با AI_DEFAULT_PROVIDER=agy مسیریابی می‌شود
This commit is contained in:
mjalmoudy94
2026-10-13 08:20:27 +03:30
parent 54074f10c3
commit 3d1e98f4fb
25 changed files with 3073 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/api/tests/test_agy.py
+715
View File
@@ -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"]