@samitouri / QOSami_Agents_Zero-1
Code Commits Issues Pulls Sessions CI/CD
main
main
py 222 lines 5.37 KB
Raw
1 from __future__ import annotations
2
3 import json
4 import os
5 import tempfile
6 import importlib
7 from dataclasses import asdict, dataclass, field
8 from pathlib import Path
9 from typing import Any, Protocol
10
11
12 CODEX_PROVIDER_ID = "codex_oauth"
13 GITHUB_COPILOT_PROVIDER_ID = "github_copilot_oauth"
14 GEMINI_API_PROVIDER_ID = "gemini_api_oauth"
15 XAI_GROK_PROVIDER_ID = "xai_grok_oauth"
16 DUMMY_API_KEY = "oauth"
17
18
19 class ProviderError(Exception):
20 def __init__(
21 self,
22 message: str,
23 *,
24 code: str = "provider_error",
25 status: int = 400,
26 details: dict[str, Any] | None = None,
27 ) -> None:
28 super().__init__(message)
29 self.message = message
30 self.code = code
31 self.status = status
32 self.details = dict(details or {})
33
34 def to_dict(self) -> dict[str, Any]:
35 result: dict[str, Any] = {
36 "ok": False,
37 "error": self.message,
38 "code": self.code,
39 "status": self.status,
40 }
41 if self.details:
42 result["details"] = self.details
43 return result
44
45
46 @dataclass(frozen=True)
47 class OAuthProviderMetadata:
48 provider_id: str
49 display_name: str
50 short_name: str
51 model_provider_id: str
52 icon: str
53 auth_flow: str
54 default_model: str
55 default_models: list[str] = field(default_factory=list)
56 proxy_base_path: str = ""
57 callback_path: str = ""
58 supports_manual_callback: bool = False
59 supports_enterprise_domain: bool = False
60 supports_oauth_client_config: bool = False
61 supports_quota_project: bool = False
62 note: str = ""
63 warning: str = ""
64
65 def to_dict(self) -> dict[str, Any]:
66 return asdict(self)
67
68
69 @dataclass(frozen=True)
70 class LoginStartResult:
71 ok: bool
72 provider_id: str
73 flow: str
74 auth_url: str = ""
75 redirect_uri: str = ""
76 attempt_id: str = ""
77 verification_url: str = ""
78 user_code: str = ""
79 interval: int = 5
80 expires_at: float = 0
81 message: str = ""
82 error: str = ""
83
84 def to_dict(self) -> dict[str, Any]:
85 return asdict(self)
86
87
88 @dataclass(frozen=True)
89 class LoginPollResult:
90 ok: bool
91 provider_id: str
92 completed: bool = False
93 account_label: str = ""
94 account_id: str = ""
95 interval: int = 5
96 expires_at: float = 0
97 expired: bool = False
98 error: str = ""
99 warning: str = ""
100
101 def to_dict(self) -> dict[str, Any]:
102 return asdict(self)
103
104
105 @dataclass(frozen=True)
106 class CallbackResult:
107 ok: bool
108 provider_id: str
109 account_label: str = ""
110 account_id: str = ""
111 error: str = ""
112
113 def to_dict(self) -> dict[str, Any]:
114 return asdict(self)
115
116
117 class OAuthProvider(Protocol):
118 provider_id: str
119
120 def metadata(self) -> OAuthProviderMetadata:
121 ...
122
123 def status(self) -> dict[str, Any]:
124 ...
125
126 def start_login(self, input: dict[str, Any] | None = None, request: Any = None) -> LoginStartResult:
127 ...
128
129 def poll_login(self, input: dict[str, Any] | None = None, request: Any = None) -> LoginPollResult:
130 ...
131
132 def complete_callback(self, args: dict[str, Any], request: Any) -> CallbackResult:
133 ...
134
135 def manual_callback(self, input: dict[str, Any], request: Any) -> LoginPollResult:
136 ...
137
138 def models(self) -> list[str]:
139 ...
140
141 def disconnect(self) -> dict[str, Any]:
142 ...
143
144 def api_key(self) -> str:
145 ...
146
147 def register_routes(self, app: Any) -> None:
148 ...
149
150
151 def provider_data_dir(provider_slug: str) -> Path:
152 if not _valid_provider_slug(provider_slug):
153 raise ProviderError(
154 "Invalid OAuth provider storage slug.",
155 code="invalid_provider_slug",
156 )
157
158 files = importlib.import_module("helpers.files")
159
160 path = Path(
161 files.get_abs_path(
162 files.USER_DIR,
163 files.PLUGINS_DIR,
164 "_oauth",
165 provider_slug,
166 )
167 )
168 path.mkdir(parents=True, exist_ok=True)
169 return path
170
171
172 def provider_auth_path(provider_slug: str) -> Path:
173 return provider_data_dir(provider_slug) / "auth.json"
174
175
176 def _valid_provider_slug(provider_slug: str) -> bool:
177 if not isinstance(provider_slug, str):
178 return False
179 if not provider_slug or provider_slug in {".", ".."}:
180 return False
181 if "/" in provider_slug or "\\" in provider_slug:
182 return False
183 return all(char.isalnum() or char in {"_", "-"} for char in provider_slug)
184
185
186 def read_json_file(path: Path) -> dict[str, Any]:
187 try:
188 with path.open("r", encoding="utf-8") as handle:
189 payload = json.load(handle)
190 except FileNotFoundError:
191 return {}
192 if not isinstance(payload, dict):
193 return {}
194 return payload
195
196
197 def write_private_json(path: Path, data: dict[str, Any]) -> None:
198 path.parent.mkdir(parents=True, exist_ok=True)
199 tmp_name = ""
200 try:
201 fd, tmp_name = tempfile.mkstemp(
202 prefix=f".{path.name}.",
203 suffix=".tmp",
204 dir=str(path.parent),
205 )
206 os.chmod(tmp_name, 0o600)
207 with os.fdopen(fd, "w", encoding="utf-8") as handle:
208 json.dump(data, handle, indent=2)
209 handle.write("\n")
210 os.replace(tmp_name, path)
211 except Exception:
212 if tmp_name:
213 Path(tmp_name).unlink(missing_ok=True)
214 raise
215 try:
216 path.chmod(0o600)
217 except OSError:
218 pass
219
220
221 def public_error(exc: Exception) -> str:
222 return str(exc) or exc.__class__.__name__