main
py 665 lines 23.6 KB
Raw
1 from __future__ import annotations
2
3 import base64
4 import json
5 import time
6 from pathlib import Path
7 from typing import Any
8 from urllib.parse import urlencode, urlparse
9
10 from plugins._oauth.helpers.providers.base import (
11 DUMMY_API_KEY,
12 GEMINI_API_PROVIDER_ID,
13 CallbackResult,
14 LoginPollResult,
15 LoginStartResult,
16 OAuthProviderMetadata,
17 ProviderError,
18 provider_auth_path,
19 read_json_file,
20 write_private_json,
21 )
22 from plugins._oauth.helpers.providers.common import (
23 as_int as _as_int,
24 as_optional_string as _as_optional_string,
25 error_message as _error_message,
26 expires_ms as _expires_ms,
27 json_payload as _json_payload,
28 latest_attempt,
29 models_from_payload as _models_from_payload,
30 parse_manual_callback,
31 )
32 from plugins._oauth.helpers.state import get_attempt, pop_attempt, put_attempt
33
34
35 GOOGLE_AUTHORIZATION_ENDPOINT = "https://accounts.google.com/o/oauth2/v2/auth"
36 GOOGLE_TOKEN_ENDPOINT = "https://oauth2.googleapis.com/token"
37 GEMINI_OPENAI_API_BASE = "https://generativelanguage.googleapis.com/v1beta/openai"
38 CURATED_MODELS = [
39 "gemini-3.5-flash",
40 "gemini-3.1-pro-preview",
41 "gemini-3-flash-preview",
42 "gemini-3.1-flash-lite",
43 "gemini-2.5-pro",
44 "gemini-2.5-flash",
45 "gemini-2.5-flash-lite",
46 ]
47 NOT_CONNECTED_MESSAGE = "Google Gemini API OAuth is not connected yet."
48 CLIENT_CONFIG_NOTE = (
49 "Requires a Google Cloud OAuth client with the Generative Language API enabled. "
50 "This uses Gemini API billing/quotas, not Antigravity or Gemini Code Assist subscription quota."
51 )
52 REFRESH_MARGIN_MS = 60_000
53
54
55 class GeminiApiOAuthProvider:
56 provider_id = GEMINI_API_PROVIDER_ID
57
58 def auth_path(self) -> Path:
59 return provider_auth_path("gemini_api")
60
61 def read_auth(self) -> dict[str, Any]:
62 return read_json_file(self.auth_path())
63
64 def write_auth(self, data: dict[str, Any]) -> None:
65 write_private_json(self.auth_path(), data)
66
67 def metadata(self) -> OAuthProviderMetadata:
68 cfg = _gemini_api_config()
69 base_path = cfg["proxy_base_path"]
70 return OAuthProviderMetadata(
71 provider_id=GEMINI_API_PROVIDER_ID,
72 display_name="Google Cloud Gemini",
73 short_name="Google Cloud",
74 model_provider_id=GEMINI_API_PROVIDER_ID,
75 icon="google",
76 auth_flow="browser_pkce",
77 default_model="gemini-3.5-flash",
78 default_models=list(CURATED_MODELS),
79 proxy_base_path=base_path,
80 callback_path=cfg["callback_path"],
81 supports_manual_callback=True,
82 supports_oauth_client_config=True,
83 supports_quota_project=True,
84 note=CLIENT_CONFIG_NOTE,
85 )
86
87 def status(self) -> dict[str, Any]:
88 cfg = _gemini_api_config()
89 auth = self.read_auth()
90 access = str(auth.get("access") or "")
91 refresh = str(auth.get("refresh") or "")
92 client_id = str(cfg.get("client_id") or auth.get("client_id") or "")
93 quota_project_id = str(cfg.get("quota_project_id") or auth.get("quota_project_id") or "")
94 result = {
95 **self.metadata().to_dict(),
96 "enabled": cfg["enabled"],
97 "connected": bool(access and refresh),
98 "account_label": _account_label(auth) if access or refresh else "",
99 "client_id": client_id,
100 "client_secret_configured": bool(cfg.get("client_secret") or auth.get("client_secret")),
101 "quota_project_id": quota_project_id,
102 "base_url": safe_api_base_url(auth.get("base_url") or cfg.get("api_base_url")),
103 "auth_file_path": str(self.auth_path()),
104 "v1_base_path": f'{cfg["proxy_base_path"]}/v1',
105 }
106 if access and refresh and _as_int(auth.get("expires"), 0) <= int(time.time() * 1000):
107 result["warning"] = "Google Gemini API OAuth access token is expired and will be refreshed on the next request."
108 return result
109
110 def start_login(self, input: dict[str, Any] | None = None, request: Any = None) -> LoginStartResult:
111 data = input or {}
112 cfg = _gemini_api_config()
113 if not cfg["enabled"]:
114 return LoginStartResult(
115 ok=False,
116 provider_id=GEMINI_API_PROVIDER_ID,
117 flow="browser_pkce",
118 error="Google Gemini API OAuth connection is disabled.",
119 )
120
121 try:
122 client = _client_config(data, cfg)
123 codex = _codex_helper()
124 pkce = codex.generate_pkce()
125 state = codex.generate_state()
126 redirect_uri = _redirect_uri(request, cfg["callback_path"])
127 attempt = put_attempt(
128 state,
129 pkce.verifier,
130 redirect_uri,
131 provider_id=GEMINI_API_PROVIDER_ID,
132 extra={
133 "client_id": client["client_id"],
134 "client_secret": client["client_secret"],
135 "quota_project_id": client["quota_project_id"],
136 "scope": " ".join(cfg["scopes"]),
137 "code_challenge": pkce.challenge,
138 "token_endpoint": GOOGLE_TOKEN_ENDPOINT,
139 "api_base_url": safe_api_base_url(cfg["api_base_url"]),
140 },
141 )
142 query = {
143 "response_type": "code",
144 "client_id": client["client_id"],
145 "redirect_uri": redirect_uri,
146 "scope": " ".join(cfg["scopes"]),
147 "code_challenge": pkce.challenge,
148 "code_challenge_method": "S256",
149 "state": state,
150 "access_type": "offline",
151 "prompt": "consent",
152 "include_granted_scopes": "true",
153 }
154 auth_url = f"{GOOGLE_AUTHORIZATION_ENDPOINT}?{urlencode(query)}"
155 except Exception as exc:
156 return LoginStartResult(
157 ok=False,
158 provider_id=GEMINI_API_PROVIDER_ID,
159 flow="browser_pkce",
160 error=str(exc),
161 message=str(exc),
162 )
163
164 return LoginStartResult(
165 ok=True,
166 provider_id=GEMINI_API_PROVIDER_ID,
167 flow="browser_pkce",
168 auth_url=auth_url,
169 redirect_uri=redirect_uri,
170 expires_at=attempt.expires_at,
171 )
172
173 def poll_login(self, input: dict[str, Any] | None = None, request: Any = None) -> LoginPollResult:
174 del input, request
175 return LoginPollResult(
176 ok=False,
177 provider_id=GEMINI_API_PROVIDER_ID,
178 error="Google Gemini API OAuth uses browser callback login.",
179 )
180
181 def exchange_code(
182 self,
183 token_endpoint: str,
184 code: str,
185 redirect_uri: str,
186 code_verifier: str,
187 client_id: str,
188 client_secret: str,
189 ) -> dict[str, Any]:
190 import requests
191
192 _validate_google_token_endpoint(token_endpoint)
193 data = {
194 "grant_type": "authorization_code",
195 "code": code,
196 "redirect_uri": redirect_uri,
197 "client_id": client_id,
198 "code_verifier": code_verifier,
199 }
200 if client_secret:
201 data["client_secret"] = client_secret
202
203 response = requests.post(
204 token_endpoint,
205 headers={
206 "Accept": "application/json",
207 "Content-Type": "application/x-www-form-urlencoded",
208 },
209 data=data,
210 timeout=30,
211 )
212 payload = _json_payload(response)
213 if not response.ok:
214 raise ProviderError(
215 _error_message(payload, f"Google Gemini API token exchange failed with status {response.status_code}."),
216 code="token_exchange_failed",
217 status=response.status_code,
218 )
219 _validate_token_payload(payload, require_refresh=True)
220 return payload
221
222 def manual_callback(self, input: dict[str, Any], request: Any = None) -> LoginPollResult:
223 del request
224 raw = input.get("callback")
225 if raw is None:
226 raw = input.get("callback_url")
227 return self._complete_from_callback(parse_manual_callback(raw), allow_missing_state=True)
228
229 def complete_callback(
230 self,
231 args: dict[str, Any],
232 request: Any = None,
233 ) -> CallbackResult:
234 del request
235 callback = {
236 "code": _as_optional_string(args.get("code")),
237 "state": _as_optional_string(args.get("state")),
238 "error": _as_optional_string(args.get("error")),
239 "error_description": _as_optional_string(args.get("error_description")),
240 }
241 result = self._complete_from_callback(callback, allow_missing_state=False)
242 return CallbackResult(
243 ok=result.ok,
244 provider_id=GEMINI_API_PROVIDER_ID,
245 account_label=result.account_label,
246 account_id=result.account_id,
247 error=result.error,
248 )
249
250 def _complete_from_callback(
251 self,
252 callback: dict[str, str | None] | None,
253 *,
254 allow_missing_state: bool,
255 ) -> LoginPollResult:
256 if not callback:
257 return LoginPollResult(ok=False, provider_id=GEMINI_API_PROVIDER_ID, error="Missing OAuth callback.")
258 if callback.get("error"):
259 return LoginPollResult(
260 ok=False,
261 provider_id=GEMINI_API_PROVIDER_ID,
262 error=str(callback.get("error_description") or callback.get("error")),
263 )
264
265 code = str(callback.get("code") or "").strip()
266 if not code:
267 return LoginPollResult(
268 ok=False,
269 provider_id=GEMINI_API_PROVIDER_ID,
270 error="The OAuth callback did not include an authorization code.",
271 )
272
273 state = str(callback.get("state") or "").strip()
274 attempt = None
275 if state:
276 attempt = get_attempt(state)
277 if attempt is None:
278 if latest_attempt(GEMINI_API_PROVIDER_ID) is not None:
279 return LoginPollResult(
280 ok=False,
281 provider_id=GEMINI_API_PROVIDER_ID,
282 error="OAuth state mismatch. Return to Agent Zero and start a new Google Gemini API connection.",
283 )
284 return LoginPollResult(
285 ok=False,
286 provider_id=GEMINI_API_PROVIDER_ID,
287 expired=True,
288 error="OAuth sign-in expired. Return to Agent Zero and start a new Google Gemini API connection.",
289 )
290 if attempt.provider_id != GEMINI_API_PROVIDER_ID:
291 return LoginPollResult(
292 ok=False,
293 provider_id=GEMINI_API_PROVIDER_ID,
294 error="OAuth state mismatch. Return to Agent Zero and start a new Google Gemini API connection.",
295 )
296 elif allow_missing_state:
297 attempt = latest_attempt(GEMINI_API_PROVIDER_ID)
298 if attempt is None:
299 return LoginPollResult(
300 ok=False,
301 provider_id=GEMINI_API_PROVIDER_ID,
302 error="No active Google Gemini API sign-in attempt was found.",
303 )
304 state = attempt.state
305 else:
306 return LoginPollResult(
307 ok=False,
308 provider_id=GEMINI_API_PROVIDER_ID,
309 error="The OAuth callback did not include state.",
310 )
311
312 token_endpoint = str(attempt.extra.get("token_endpoint") or GOOGLE_TOKEN_ENDPOINT)
313 client_id = str(attempt.extra.get("client_id") or "")
314 client_secret = str(attempt.extra.get("client_secret") or "")
315 try:
316 payload = self.exchange_code(
317 token_endpoint,
318 code,
319 attempt.redirect_uri,
320 attempt.verifier,
321 client_id,
322 client_secret,
323 )
324 auth = _auth_from_token_payload(
325 payload,
326 token_endpoint,
327 client_id=client_id,
328 client_secret=client_secret,
329 quota_project_id=str(attempt.extra.get("quota_project_id") or ""),
330 api_base_url=str(attempt.extra.get("api_base_url") or GEMINI_OPENAI_API_BASE),
331 )
332 self.write_auth(auth)
333 pop_attempt(state)
334 except Exception as exc:
335 return LoginPollResult(
336 ok=False,
337 provider_id=GEMINI_API_PROVIDER_ID,
338 error=str(exc),
339 )
340
341 label = _account_label(auth)
342 return LoginPollResult(
343 ok=True,
344 provider_id=GEMINI_API_PROVIDER_ID,
345 completed=True,
346 account_label=label,
347 account_id=label,
348 )
349
350 def ensure_fresh_auth(self) -> dict[str, Any]:
351 auth = self.read_auth()
352 refresh = str(auth.get("refresh") or "")
353 if not refresh:
354 return auth
355
356 access = str(auth.get("access") or "")
357 expires = _as_int(auth.get("expires"), 0)
358 if access and expires and expires > int(time.time() * 1000) + REFRESH_MARGIN_MS:
359 return auth
360
361 token_endpoint = str(auth.get("token_endpoint") or GOOGLE_TOKEN_ENDPOINT)
362 _validate_google_token_endpoint(token_endpoint)
363 client_id = str(auth.get("client_id") or "")
364 client_secret = str(auth.get("client_secret") or "")
365 cfg = _gemini_api_config()
366 if cfg.get("client_id"):
367 client_id = str(cfg["client_id"])
368 if cfg.get("client_secret"):
369 client_secret = str(cfg["client_secret"])
370 if not client_id:
371 raise ProviderError(
372 "Google Gemini API OAuth refresh requires the original OAuth client ID.",
373 code="missing_oauth_client",
374 status=401,
375 )
376
377 refreshed = self._refresh_tokens(token_endpoint, refresh, client_id, client_secret, auth)
378 self.write_auth(refreshed)
379 return refreshed
380
381 def _refresh_tokens(
382 self,
383 token_endpoint: str,
384 refresh: str,
385 client_id: str,
386 client_secret: str,
387 existing: dict[str, Any],
388 ) -> dict[str, Any]:
389 import requests
390
391 data = {
392 "grant_type": "refresh_token",
393 "refresh_token": refresh,
394 "client_id": client_id,
395 }
396 if client_secret:
397 data["client_secret"] = client_secret
398 response = requests.post(
399 token_endpoint,
400 headers={
401 "Accept": "application/json",
402 "Content-Type": "application/x-www-form-urlencoded",
403 },
404 data=data,
405 timeout=30,
406 )
407 payload = _json_payload(response)
408 if not response.ok:
409 raise ProviderError(
410 _error_message(payload, f"Google Gemini API token refresh failed with status {response.status_code}."),
411 code="token_refresh_failed",
412 status=response.status_code,
413 )
414 _validate_token_payload(payload, require_refresh=False)
415 merged = dict(existing)
416 merged.update(
417 _auth_from_token_payload(
418 payload,
419 token_endpoint,
420 fallback_refresh=refresh,
421 client_id=client_id,
422 client_secret=client_secret,
423 quota_project_id=str(existing.get("quota_project_id") or ""),
424 api_base_url=str(existing.get("base_url") or GEMINI_OPENAI_API_BASE),
425 )
426 )
427 if not payload.get("id_token") and existing.get("id_token"):
428 merged["id_token"] = existing["id_token"]
429 return merged
430
431 def models(self) -> list[str]:
432 if not self.read_auth():
433 return list(CURATED_MODELS)
434 try:
435 auth = self.ensure_fresh_auth()
436 except Exception:
437 return list(CURATED_MODELS)
438 access = str(auth.get("access") or "")
439 if not access:
440 return list(CURATED_MODELS)
441
442 base_url = safe_api_base_url(auth.get("base_url"))
443 try:
444 import requests
445
446 response = requests.get(
447 f"{base_url}/models",
448 headers=_gemini_headers(auth),
449 timeout=30,
450 )
451 if not response.ok:
452 return list(CURATED_MODELS)
453 parsed = _models_from_payload(response.json())
454 return parsed or list(CURATED_MODELS)
455 except Exception:
456 return list(CURATED_MODELS)
457
458 def disconnect(self) -> dict[str, Any]:
459 path = self.auth_path()
460 existed = path.exists()
461 try:
462 path.unlink(missing_ok=True)
463 except FileNotFoundError:
464 pass
465 return {
466 "disconnected": existed,
467 "removed_auth_files": [str(path)] if existed else [],
468 }
469
470 def api_key(self) -> str:
471 return DUMMY_API_KEY
472
473 def register_routes(self, app: Any) -> None:
474 from plugins._oauth.helpers import routes
475
476 route_defs = [
477 ("/oauth/gemini-api/health", "oauth_gemini_api_health", routes.gemini_api_health, ["GET"]),
478 ("/oauth/gemini-api/callback", "oauth_gemini_api_callback", routes.gemini_api_callback, ["GET"]),
479 (
480 "/oauth/gemini-api/v1/models",
481 "oauth_gemini_api_models",
482 routes.gemini_api_models,
483 ["GET", "OPTIONS"],
484 ),
485 (
486 "/oauth/gemini-api/v1/chat/completions",
487 "oauth_gemini_api_chat_completions",
488 routes.gemini_api_chat_completions,
489 ["POST", "OPTIONS"],
490 ),
491 (
492 "/oauth/gemini-api/v1/responses",
493 "oauth_gemini_api_responses",
494 routes.gemini_api_responses,
495 ["POST", "OPTIONS"],
496 ),
497 ]
498 for rule, endpoint, view_func, methods in route_defs:
499 if endpoint in app.view_functions:
500 continue
501 app.add_url_rule(rule, endpoint, view_func, methods=methods)
502
503
504 def _auth_from_token_payload(
505 payload: dict[str, Any],
506 token_endpoint: str,
507 *,
508 fallback_refresh: str = "",
509 client_id: str = "",
510 client_secret: str = "",
511 quota_project_id: str = "",
512 api_base_url: str = GEMINI_OPENAI_API_BASE,
513 ) -> dict[str, Any]:
514 return {
515 "provider": GEMINI_API_PROVIDER_ID,
516 "type": "oauth",
517 "access": str(payload.get("access_token") or ""),
518 "refresh": str(payload.get("refresh_token") or fallback_refresh or ""),
519 "expires": _expires_ms(payload),
520 "id_token": str(payload.get("id_token") or ""),
521 "token_type": str(payload.get("token_type") or "Bearer"),
522 "scope": str(payload.get("scope") or ""),
523 "token_endpoint": token_endpoint,
524 "base_url": safe_api_base_url(api_base_url),
525 "client_id": client_id,
526 "client_secret": client_secret,
527 "quota_project_id": quota_project_id,
528 }
529
530
531 def _validate_token_payload(payload: dict[str, Any], *, require_refresh: bool) -> None:
532 if not isinstance(payload, dict):
533 raise ProviderError("Google Gemini API token endpoint returned a malformed response.", code="token_malformed", status=502)
534 missing = []
535 if not str(payload.get("access_token") or ""):
536 missing.append("access_token")
537 if require_refresh and not str(payload.get("refresh_token") or ""):
538 missing.append("refresh_token")
539 if missing:
540 raise ProviderError(
541 f"Google Gemini API token response is missing: {', '.join(missing)}",
542 code="token_malformed",
543 status=502,
544 )
545
546
547 def safe_api_base_url(value: Any) -> str:
548 text = str(value or "").strip().rstrip("/")
549 if not text:
550 return GEMINI_OPENAI_API_BASE
551 parsed = urlparse(text)
552 host = (parsed.hostname or "").lower()
553 if parsed.scheme == "https" and host == "generativelanguage.googleapis.com":
554 return text
555 return GEMINI_OPENAI_API_BASE
556
557
558 def _gemini_headers(auth: dict[str, Any]) -> dict[str, str]:
559 headers = {
560 "Accept": "application/json",
561 "Authorization": f'Bearer {str(auth.get("access") or "")}',
562 }
563 quota_project_id = str(auth.get("quota_project_id") or "").strip()
564 if quota_project_id:
565 headers["x-goog-user-project"] = quota_project_id
566 return headers
567
568
569 def _client_config(data: dict[str, Any], cfg: dict[str, Any]) -> dict[str, str]:
570 client_id = str(data.get("client_id") or cfg.get("client_id") or "").strip()
571 client_secret = str(data.get("client_secret") or cfg.get("client_secret") or "").strip()
572 quota_project_id = str(data.get("quota_project_id") or cfg.get("quota_project_id") or "").strip()
573 if not client_id or not client_secret:
574 raise ProviderError(
575 "Configure a Google OAuth client ID and client secret before connecting Gemini API OAuth.",
576 code="missing_oauth_client",
577 )
578 return {
579 "client_id": client_id,
580 "client_secret": client_secret,
581 "quota_project_id": quota_project_id,
582 }
583
584
585 def _redirect_uri(request: Any, callback_path: str) -> str:
586 origin = ""
587 if request is not None:
588 origin = (getattr(request, "headers", {}).get("Origin") or "").rstrip("/")
589 if not _is_local_origin(origin):
590 origin = getattr(request, "url_root", "").rstrip("/")
591 return f"{origin}{callback_path}"
592
593
594 def _is_local_origin(origin: str) -> bool:
595 if not origin:
596 return False
597 return (
598 origin.startswith("http://localhost:")
599 or origin == "http://localhost"
600 or origin.startswith("http://127.0.0.1:")
601 or origin == "http://127.0.0.1"
602 or origin.startswith("http://[::1]:")
603 or origin == "http://[::1]"
604 )
605
606
607 def _validate_google_token_endpoint(value: str) -> None:
608 parsed = urlparse(value)
609 host = (parsed.hostname or "").lower()
610 if parsed.scheme != "https" or host != "oauth2.googleapis.com":
611 raise ProviderError(
612 "Google Gemini API OAuth token endpoint is invalid.",
613 code="invalid_token_endpoint",
614 status=502,
615 )
616
617
618 def _account_label(auth: dict[str, Any]) -> str:
619 claims = _jwt_claims(str(auth.get("id_token") or ""))
620 return str(claims.get("email") or auth.get("account_label") or "Google Gemini API")
621
622
623 def _jwt_claims(token: str) -> dict[str, Any]:
624 parts = token.split(".")
625 if len(parts) < 2:
626 return {}
627 try:
628 payload = parts[1] + "=" * (-len(parts[1]) % 4)
629 decoded = base64.urlsafe_b64decode(payload.encode("utf-8"))
630 parsed = json.loads(decoded.decode("utf-8"))
631 except Exception:
632 return {}
633 return parsed if isinstance(parsed, dict) else {}
634
635
636 def _codex_helper():
637 import importlib
638
639 return importlib.import_module("plugins._oauth.helpers.codex")
640
641
642 def _gemini_api_config() -> dict[str, Any]:
643 try:
644 from plugins._oauth.helpers.config import gemini_api_config
645
646 return gemini_api_config()
647 except ModuleNotFoundError as exc:
648 if exc.name and exc.name.startswith("plugins._oauth"):
649 raise
650 return {
651 "enabled": True,
652 "client_id": "",
653 "client_secret": "",
654 "scopes": [
655 "openid",
656 "email",
657 "profile",
658 "https://www.googleapis.com/auth/cloud-platform",
659 "https://www.googleapis.com/auth/generative-language.retriever",
660 ],
661 "quota_project_id": "",
662 "api_base_url": GEMINI_OPENAI_API_BASE,
663 "proxy_base_path": "/oauth/gemini-api",
664 "callback_path": "/oauth/gemini-api/callback",
665 }