main
py 657 lines 19.2 KB
Raw
1 from __future__ import annotations
2
3 import ipaddress
4 import json
5 import time
6 from typing import Any, Callable
7
8 from flask import Response, jsonify, request, stream_with_context
9
10 from plugins._oauth.helpers import codex
11 from plugins._oauth.helpers.config import codex_config
12 from plugins._oauth.helpers.providers import (
13 GEMINI_API_PROVIDER_ID,
14 GITHUB_COPILOT_PROVIDER_ID,
15 XAI_GROK_PROVIDER_ID,
16 ProviderError,
17 get_provider,
18 provider_registry,
19 )
20 from plugins._oauth.helpers.state import pop_attempt
21
22
23 def register_oauth_routes(app) -> None:
24 cfg = codex_config()
25 base = cfg["proxy_base_path"]
26
27 routes = [
28 (f"{base}/health", "oauth_codex_health", codex_health, ["GET"]),
29 (f"{base}/callback", "oauth_codex_callback", codex_callback, ["GET"]),
30 (cfg["callback_path"], "oauth_codex_compat_callback", codex_callback, ["GET"]),
31 (f"{base}/v1/models", "oauth_codex_models", codex_models, ["GET", "OPTIONS"]),
32 (
33 f"{base}/v1/responses",
34 "oauth_codex_responses",
35 codex_responses,
36 ["POST", "OPTIONS"],
37 ),
38 (
39 f"{base}/v1/chat/completions",
40 "oauth_codex_chat_completions",
41 codex_chat_completions,
42 ["POST", "OPTIONS"],
43 ),
44 ]
45 for rule, endpoint, view_func, methods in routes:
46 if endpoint in app.view_functions:
47 continue
48 app.add_url_rule(rule, endpoint, view_func, methods=methods)
49
50 for provider in provider_registry().values():
51 provider.register_routes(app)
52
53
54 def codex_health():
55 return jsonify({"ok": True, "provider": "codex", "base_path": codex_config()["proxy_base_path"]})
56
57
58 def codex_callback():
59 error = request.args.get("error")
60 if error:
61 description = request.args.get("error_description") or error
62 return _html_page("Codex Sign-In Failed", description), 400
63
64 state = request.args.get("state", "")
65 code = request.args.get("code", "")
66 attempt = pop_attempt(state)
67 if not attempt:
68 return _html_page("Codex Sign-In Expired", "Return to Agent Zero and start a new Codex connection."), 400
69 if not code:
70 return _html_page("Codex Sign-In Failed", "The OAuth callback did not include an authorization code."), 400
71
72 try:
73 auth = codex.complete_login(code, attempt.redirect_uri, attempt.verifier)
74 info = codex.status()
75 except Exception as exc:
76 return _html_page("Codex Sign-In Failed", str(exc)), 500
77
78 email = info.get("email") or "Connected"
79 detail = f"{email}\n{auth.account_id}"
80 return _html_page("Codex Connected", detail)
81
82
83 def codex_models():
84 if request.method == "OPTIONS":
85 return _options_response()
86 denied = _proxy_denied_response()
87 if denied:
88 return denied
89 try:
90 models = codex.fetch_models()
91 return jsonify(
92 {
93 "object": "list",
94 "data": [
95 {
96 "id": model,
97 "object": "model",
98 "created": 0,
99 "owned_by": "codex-oauth",
100 }
101 for model in models
102 ],
103 }
104 )
105 except ProviderError as exc:
106 return _json_error(str(exc), status=exc.status, code=exc.code)
107 except Exception as exc:
108 return _json_error(str(exc), status=502, code="upstream_error")
109
110
111 def codex_responses():
112 if request.method == "OPTIONS":
113 return _options_response()
114 denied = _proxy_denied_response()
115 if denied:
116 return denied
117
118 body = request.get_json(silent=True)
119 if not isinstance(body, dict):
120 return _json_error("Request body must be a JSON object.")
121
122 wants_stream = body.get("stream") is True
123 upstream_body = codex.prepare_responses_body(body, force_stream=True)
124 try:
125 upstream = codex.request_codex(
126 "/responses",
127 method="POST",
128 headers={"Content-Type": "application/json"},
129 body=json.dumps(upstream_body),
130 stream=True,
131 )
132 except Exception as exc:
133 return _json_error(str(exc), status=502, code="upstream_error")
134
135 if not upstream.ok:
136 return _copy_upstream_response(upstream)
137 if wants_stream:
138 return _stream_upstream_sse(upstream)
139
140 try:
141 completed = codex.collect_completed_response(upstream)
142 except Exception as exc:
143 return _json_error(str(exc), status=502, code="upstream_error")
144 return jsonify(completed)
145
146
147 def codex_chat_completions():
148 if request.method == "OPTIONS":
149 return _options_response()
150 denied = _proxy_denied_response()
151 if denied:
152 return denied
153
154 body = request.get_json(silent=True)
155 if not isinstance(body, dict):
156 return _json_error("Request body must be a JSON object.")
157
158 try:
159 response_body = codex.chat_messages_to_response_body(body)
160 except Exception as exc:
161 return _json_error(str(exc))
162
163 wants_stream = body.get("stream") is True
164 response_body["stream"] = True
165 try:
166 upstream = codex.request_codex(
167 "/responses",
168 method="POST",
169 headers={"Content-Type": "application/json"},
170 body=json.dumps(codex.prepare_responses_body(response_body, force_stream=True)),
171 stream=True,
172 )
173 except Exception as exc:
174 return _json_error(str(exc), status=502, code="upstream_error")
175
176 if not upstream.ok:
177 return _copy_upstream_response(upstream)
178 if wants_stream:
179 return _stream_chat_completion(upstream, str(body.get("model") or response_body["model"]))
180
181 try:
182 completed = codex.collect_completed_response(upstream)
183 except Exception as exc:
184 return _json_error(str(exc), status=502, code="upstream_error")
185
186 text = codex.response_text(completed)
187 return jsonify(
188 {
189 "id": f"chatcmpl_{int(time.time() * 1000)}",
190 "object": "chat.completion",
191 "created": int(time.time()),
192 "model": body.get("model") or response_body["model"],
193 "choices": [
194 {
195 "index": 0,
196 "message": {"role": "assistant", "content": text},
197 "finish_reason": "stop",
198 }
199 ],
200 "usage": completed.get("usage") or {},
201 }
202 )
203
204
205 def github_copilot_health():
206 return jsonify(
207 {
208 "ok": True,
209 "provider": GITHUB_COPILOT_PROVIDER_ID,
210 "base_path": "/oauth/github-copilot",
211 }
212 )
213
214
215 def github_copilot_models():
216 if request.method == "OPTIONS":
217 return _options_response()
218 denied = _proxy_denied_response()
219 if denied:
220 return denied
221
222 provider = get_provider(GITHUB_COPILOT_PROVIDER_ID)
223 return jsonify(
224 {
225 "object": "list",
226 "data": [
227 {
228 "id": model,
229 "object": "model",
230 "created": 0,
231 "owned_by": "github-copilot-oauth",
232 }
233 for model in provider.models()
234 ],
235 }
236 )
237
238
239 def github_copilot_chat_completions():
240 if request.method == "OPTIONS":
241 return _options_response()
242 return _github_copilot_json_proxy("/chat/completions")
243
244
245 def github_copilot_responses():
246 if request.method == "OPTIONS":
247 return _options_response()
248 return _github_copilot_json_proxy("/responses")
249
250
251 def xai_grok_health():
252 return jsonify(
253 {
254 "ok": True,
255 "provider": XAI_GROK_PROVIDER_ID,
256 "base_path": "/oauth/xai-grok",
257 }
258 )
259
260
261 def xai_grok_callback():
262 provider = get_provider(XAI_GROK_PROVIDER_ID)
263 result = provider.complete_callback(dict(request.args), request)
264 if result.ok:
265 return _html_page("xAI Grok Connected", result.account_label or "Connected")
266 return _html_page("xAI Grok Sign-In Failed", result.error or "The OAuth callback failed."), 400
267
268
269 def xai_grok_models():
270 if request.method == "OPTIONS":
271 return _options_response()
272 denied = _proxy_denied_response()
273 if denied:
274 return denied
275
276 provider = get_provider(XAI_GROK_PROVIDER_ID)
277 return jsonify(
278 {
279 "object": "list",
280 "data": [
281 {
282 "id": model,
283 "object": "model",
284 "created": 0,
285 "owned_by": "xai-grok-oauth",
286 }
287 for model in provider.models()
288 ],
289 }
290 )
291
292
293 def xai_grok_chat_completions():
294 if request.method == "OPTIONS":
295 return _options_response()
296 return _xai_grok_json_proxy("/chat/completions")
297
298
299 def xai_grok_responses():
300 if request.method == "OPTIONS":
301 return _options_response()
302 return _xai_grok_json_proxy("/responses")
303
304
305 def gemini_api_health():
306 return jsonify(
307 {
308 "ok": True,
309 "provider": GEMINI_API_PROVIDER_ID,
310 "base_path": "/oauth/gemini-api",
311 }
312 )
313
314
315 def gemini_api_callback():
316 provider = get_provider(GEMINI_API_PROVIDER_ID)
317 result = provider.complete_callback(dict(request.args), request)
318 if result.ok:
319 return _html_page("Google Gemini API Connected", result.account_label or "Connected")
320 return _html_page("Google Gemini API Sign-In Failed", result.error or "The OAuth callback failed."), 400
321
322
323 def gemini_api_models():
324 if request.method == "OPTIONS":
325 return _options_response()
326 denied = _proxy_denied_response()
327 if denied:
328 return denied
329
330 provider = get_provider(GEMINI_API_PROVIDER_ID)
331 return jsonify(
332 {
333 "object": "list",
334 "data": [
335 {
336 "id": model,
337 "object": "model",
338 "created": 0,
339 "owned_by": "gemini-api-oauth",
340 }
341 for model in provider.models()
342 ],
343 }
344 )
345
346
347 def gemini_api_chat_completions():
348 if request.method == "OPTIONS":
349 return _options_response()
350 return _gemini_api_json_proxy("/chat/completions")
351
352
353 def gemini_api_responses():
354 if request.method == "OPTIONS":
355 return _options_response()
356 return _gemini_api_json_proxy("/responses")
357
358
359 def _github_copilot_json_proxy(path: str):
360 from plugins._oauth.helpers.providers.github_copilot import COPILOT_HEADERS, safe_copilot_base_url
361
362 return _oauth_json_proxy(
363 GITHUB_COPILOT_PROVIDER_ID,
364 path,
365 "GitHub Copilot OAuth is not connected.",
366 lambda auth: safe_copilot_base_url(auth.get("base_url"), auth.get("enterprise_domain")),
367 lambda auth, access: {
368 **COPILOT_HEADERS,
369 "Authorization": f"Bearer {access}",
370 "Content-Type": "application/json",
371 },
372 )
373
374
375 def _xai_grok_json_proxy(path: str):
376 from plugins._oauth.helpers.providers.xai_grok import safe_api_base_url
377
378 return _oauth_json_proxy(
379 XAI_GROK_PROVIDER_ID,
380 path,
381 "xAI Grok OAuth is not connected.",
382 lambda auth: safe_api_base_url(auth.get("base_url")),
383 lambda auth, access: {
384 "Accept": "application/json",
385 "Authorization": f"Bearer {access}",
386 "Content-Type": "application/json",
387 },
388 require_refresh=True,
389 )
390
391
392 def _gemini_api_json_proxy(path: str):
393 from plugins._oauth.helpers.providers.gemini_api import _gemini_headers, safe_api_base_url
394
395 return _oauth_json_proxy(
396 GEMINI_API_PROVIDER_ID,
397 path,
398 "Google Gemini API OAuth is not connected.",
399 lambda auth: safe_api_base_url(auth.get("base_url")),
400 lambda auth, access: {
401 **_gemini_headers(auth),
402 "Content-Type": "application/json",
403 },
404 require_refresh=True,
405 )
406
407
408 def _oauth_json_proxy(
409 provider_id: str,
410 path: str,
411 not_connected_message: str,
412 base_url_for: Callable[[dict[str, Any]], str],
413 headers_for: Callable[[dict[str, Any], str], dict[str, str]],
414 *,
415 require_refresh: bool = False,
416 ):
417 denied = _proxy_denied_response()
418 if denied:
419 return denied
420
421 body = request.get_json(silent=True)
422 if not isinstance(body, dict):
423 return _json_error("Request body must be a JSON object.")
424
425 try:
426 auth = _provider_auth(provider_id)
427 except ProviderError as exc:
428 return _json_error(str(exc), status=exc.status, code=exc.code)
429 except Exception as exc:
430 return _json_error(str(exc), status=502, code="upstream_error")
431
432 access = str(auth.get("access") or "")
433 refresh = str(auth.get("refresh") or "")
434 if not access or (require_refresh and not refresh):
435 return _json_error(not_connected_message, status=401, code="not_connected")
436
437 wants_stream = body.get("stream") is True
438 try:
439 import requests
440
441 base_url = base_url_for(auth)
442 upstream = requests.post(
443 f"{base_url}{path}",
444 headers=headers_for(auth, access),
445 json=body,
446 stream=wants_stream,
447 timeout=120,
448 )
449 except ProviderError as exc:
450 return _json_error(str(exc), status=exc.status, code=exc.code)
451 except Exception as exc:
452 return _json_error(str(exc), status=502, code="upstream_error")
453
454 if wants_stream and upstream.ok:
455 return _stream_upstream_sse(upstream)
456 return _copy_upstream_response(upstream)
457
458
459 def _provider_auth(provider_id: str) -> dict[str, Any]:
460 provider = get_provider(provider_id)
461 ensure_fresh_auth = getattr(provider, "ensure_fresh_auth", None)
462 read_auth = getattr(provider, "read_auth", None)
463 if callable(ensure_fresh_auth):
464 auth = ensure_fresh_auth()
465 elif callable(read_auth):
466 auth = read_auth()
467 else:
468 auth = {}
469 return auth if isinstance(auth, dict) else {}
470
471
472 def _stream_upstream_sse(upstream):
473 headers = codex.response_headers(upstream)
474 headers.setdefault("Content-Type", "text/event-stream")
475 headers.setdefault("Cache-Control", "no-cache")
476 return Response(
477 stream_with_context(upstream.iter_content(chunk_size=8192)),
478 status=upstream.status_code,
479 headers=headers,
480 )
481
482
483 def _stream_chat_completion(upstream, model: str):
484 created = int(time.time())
485 chunk_id = f"chatcmpl_{int(time.time() * 1000)}"
486
487 def generate():
488 yield _sse_data(
489 {
490 "id": chunk_id,
491 "object": "chat.completion.chunk",
492 "created": created,
493 "model": model,
494 "choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}],
495 }
496 )
497 for event in codex.iter_sse_events(upstream):
498 data = event.get("data")
499 if not data:
500 continue
501 try:
502 parsed = json.loads(data)
503 except json.JSONDecodeError:
504 continue
505 if not isinstance(parsed, dict):
506 continue
507 for delta in codex.extract_sse_text_deltas(parsed, event.get("event", "")):
508 yield _sse_data(
509 {
510 "id": chunk_id,
511 "object": "chat.completion.chunk",
512 "created": created,
513 "model": model,
514 "choices": [
515 {
516 "index": 0,
517 "delta": {"content": delta},
518 "finish_reason": None,
519 }
520 ],
521 }
522 )
523 yield _sse_data(
524 {
525 "id": chunk_id,
526 "object": "chat.completion.chunk",
527 "created": created,
528 "model": model,
529 "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
530 }
531 )
532 yield "data: [DONE]\n\n"
533
534 return Response(
535 stream_with_context(generate()),
536 headers={"Content-Type": "text/event-stream", "Cache-Control": "no-cache"},
537 )
538
539
540 def _copy_upstream_response(upstream):
541 return Response(
542 upstream.content,
543 status=upstream.status_code,
544 headers=codex.response_headers(upstream),
545 )
546
547
548 def _proxy_denied_response() -> Response | None:
549 if _proxy_authorized():
550 return None
551 return _json_error("Codex/ChatGPT account proxy access denied.", status=403, code="access_denied")
552
553
554 def _proxy_authorized() -> bool:
555 cfg = codex_config()
556 token = cfg["proxy_token"]
557 supplied = _supplied_proxy_token()
558 if token and supplied == token:
559 return True
560 if cfg["require_proxy_token"]:
561 return False
562 return _remote_is_loopback(request.remote_addr)
563
564
565 def _supplied_proxy_token() -> str:
566 auth = request.headers.get("Authorization", "")
567 if auth.lower().startswith("bearer "):
568 return auth[7:].strip()
569 return (
570 request.headers.get("X-API-Key")
571 or request.args.get("api_key")
572 or request.args.get("key")
573 or ""
574 ).strip()
575
576
577 def _host_is_local(host: str) -> bool:
578 hostname = (host or "").strip().lower()
579 if hostname.startswith("["):
580 closing_bracket = hostname.find("]")
581 hostname = hostname[1:closing_bracket] if closing_bracket >= 0 else hostname.strip("[]")
582 elif hostname.count(":") == 1:
583 hostname = hostname.split(":", 1)[0]
584 if hostname in {"localhost", "127.0.0.1", "::1"}:
585 return True
586 try:
587 return ipaddress.ip_address(hostname).is_loopback
588 except ValueError:
589 return False
590
591
592 def _remote_is_loopback(addr: str | None) -> bool:
593 try:
594 return ipaddress.ip_address(addr or "").is_loopback
595 except ValueError:
596 return False
597
598
599 def _json_error(message: str, *, status: int = 400, code: str = "invalid_request") -> Response:
600 return jsonify({"error": {"message": message, "type": code, "code": code}}), status
601
602
603 def _options_response() -> Response:
604 return Response(status=204)
605
606
607 def _sse_data(payload: dict[str, Any]) -> str:
608 return f"data: {json.dumps(payload, separators=(',', ':'))}\n\n"
609
610
611 def _html_page(title: str, body: str) -> str:
612 return f"""<!doctype html>
613 <html lang="en">
614 <head>
615 <meta charset="utf-8">
616 <meta name="viewport" content="width=device-width,initial-scale=1">
617 <title>{_escape_html(title)}</title>
618 <style>
619 body {{
620 margin: 0;
621 min-height: 100vh;
622 display: grid;
623 place-items: center;
624 background: #101214;
625 color: #f2f5f7;
626 font-family: Inter, ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif;
627 }}
628 main {{
629 width: min(560px, calc(100vw - 32px));
630 border: 1px solid rgba(255,255,255,.14);
631 border-radius: 8px;
632 padding: 24px;
633 background: #171a1d;
634 box-shadow: 0 18px 70px rgba(0,0,0,.28);
635 }}
636 h1 {{ margin: 0 0 10px; font-size: 24px; }}
637 p {{ margin: 0; color: #b9c1c9; line-height: 1.5; white-space: pre-line; }}
638 span {{ color: #7f8b96; font-size: 13px; }}
639 </style>
640 </head>
641 <body>
642 <main>
643 <h1>{_escape_html(title)}</h1>
644 <p>{_escape_html(body)}</p>
645 </main>
646 </body>
647 </html>"""
648
649
650 def _escape_html(value: str) -> str:
651 return (
652 value.replace("&", "&amp;")
653 .replace("<", "&lt;")
654 .replace(">", "&gt;")
655 .replace('"', "&quot;")
656 .replace("'", "&#39;")
657 )