Skip to content

Commit 3b6eea1

Browse files
committed
feat(xai): enable xAI Grok device-code OAuth in admin WebUI
Add x-ai to the proxy's admin OAuth API so the WebUI Credentials page can onboard xAI Grok accounts via the device-code flow. src/proxy_app/api/oauth.py: - Add 'x-ai' to PROVIDER_META with flow_type=device_code. - Add dispatcher branch in start_oauth_flow for 'x-ai'. - Implement _start_xai_device_flow, _poll_xai_device, and _finalize_xai parallel to the copilot device-code path. Reuses XAI_CLIENT_ID, XAI_OAUTH_SCOPES, XAI_DEVICE_CODE_URL, XAI_TOKEN_URL, and XAI_USERINFO_URL from XAiAuthBase — no constant duplication. - Persists credentials via the existing _save_credential_file helper with the exact shape XAiAuthBase's loader expects (access_token, refresh_token, expiry_date, account_id, _proxy_metadata). No frontend changes. The existing Credentials.tsx dialog already handles flow_type=device_code generically (user_code, verification_uri, polling, copy-to-clipboard) — adding the provider to PROVIDER_META is sufficient to surface it in the 'Add OAuth' dialog. tests/test_xai_oauth_flow.py: 6 test cases covering the providers list, the start envelope (with upstream mock), the unknown-provider regression, the status endpoint, and the credential-file prefix. Un-ignore via .gitignore per the established pattern for tracked test files in this repo.
1 parent 7e90659 commit 3b6eea1

3 files changed

Lines changed: 449 additions & 0 deletions

File tree

.gitignore

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -142,6 +142,7 @@ tests/*
142142
!tests/test_classifier_scoped_routing.py
143143
!tests/test_session_tracking.py
144144
!tests/test_selection_engine.py
145+
!tests/test_xai_oauth_flow.py
145146
!tests/test_usage_reconciliation.py
146147
docs/ignored/
147148
.env

src/proxy_app/api/oauth.py

Lines changed: 163 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,11 @@
4949
"flow_type": "device_code",
5050
"description": "GitHub Copilot via device flow. Enter code at GitHub.",
5151
},
52+
"x-ai": {
53+
"name": "xAI Grok",
54+
"flow_type": "device_code",
55+
"description": "xAI Grok via device flow. Enter code at auth.x.ai.",
56+
},
5257
}
5358

5459

@@ -109,6 +114,8 @@ async def start_oauth_flow(req: OAuthStartRequest):
109114

110115
if provider == "copilot":
111116
return await _start_copilot_device_flow(flow_id, flow)
117+
elif provider == "x-ai":
118+
return await _start_xai_device_flow(flow_id, flow)
112119
elif provider in ("codex", "gemini_cli", "anthropic"):
113120
return _start_paste_flow(flow_id, flow, provider)
114121
else:
@@ -267,6 +274,162 @@ async def _finalize_copilot(flow_id: str, flow: dict, github_token: str, client:
267274
}
268275

269276

277+
# ---------------------------------------------------------------------------
278+
# xAI Grok: device flow
279+
# ---------------------------------------------------------------------------
280+
async def _start_xai_device_flow(flow_id: str, flow: dict):
281+
# Reuse xAI provider constants — public client, no client_secret.
282+
from rotator_library.providers.x_ai_auth_base import (
283+
XAI_CLIENT_ID,
284+
XAI_OAUTH_SCOPES,
285+
XAI_DEVICE_CODE_URL,
286+
)
287+
288+
async with httpx.AsyncClient() as client:
289+
resp = await client.post(
290+
XAI_DEVICE_CODE_URL,
291+
data={
292+
"client_id": XAI_CLIENT_ID,
293+
"scope": " ".join(XAI_OAUTH_SCOPES),
294+
},
295+
headers={
296+
"Content-Type": "application/x-www-form-urlencoded",
297+
"Accept": "application/json",
298+
},
299+
timeout=30.0,
300+
)
301+
if not resp.is_success:
302+
raise HTTPException(
303+
502, f"xAI device code request failed: {resp.text}"
304+
)
305+
306+
data = resp.json()
307+
308+
flow["device_code"] = data["device_code"]
309+
flow["interval"] = data.get("interval", 5)
310+
flow["expires_in"] = data.get("expires_in", 600)
311+
flow["client_id"] = XAI_CLIENT_ID
312+
_pending_flows[flow_id] = flow
313+
314+
# Start background polling
315+
asyncio.create_task(_poll_xai_device(flow_id))
316+
317+
return {
318+
"flow_id": flow_id,
319+
"flow_type": "device_code",
320+
"verification_uri": data.get("verification_uri", "https://auth.x.ai/device"),
321+
"user_code": data.get("user_code", ""),
322+
"expires_in": data.get("expires_in", 600),
323+
}
324+
325+
326+
async def _poll_xai_device(flow_id: str):
327+
flow = _pending_flows.get(flow_id)
328+
if not flow:
329+
return
330+
331+
from rotator_library.providers.x_ai_auth_base import XAI_TOKEN_URL
332+
333+
client_id = flow["client_id"]
334+
device_code = flow["device_code"]
335+
interval = flow["interval"]
336+
max_polls = flow["expires_in"] // interval
337+
338+
async with httpx.AsyncClient() as client:
339+
for _ in range(max_polls):
340+
await asyncio.sleep(interval)
341+
if flow_id not in _pending_flows:
342+
return
343+
344+
try:
345+
resp = await client.post(
346+
XAI_TOKEN_URL,
347+
data={
348+
"client_id": client_id,
349+
"device_code": device_code,
350+
"grant_type": "urn:ietf:params:oauth:grant-type:device_code",
351+
},
352+
headers={
353+
"Content-Type": "application/x-www-form-urlencoded",
354+
"Accept": "application/json",
355+
},
356+
timeout=30.0,
357+
)
358+
if not resp.is_success:
359+
continue
360+
361+
token_data = resp.json()
362+
if "access_token" in token_data:
363+
await _finalize_xai(flow_id, flow, token_data, client)
364+
return
365+
366+
error = token_data.get("error", "")
367+
if error == "expired_token":
368+
flow["status"] = "error"
369+
flow["error"] = "Device code expired. Please try again."
370+
return
371+
except Exception as e:
372+
lib_logger.debug(f"xAI poll error: {e}")
373+
continue
374+
375+
flow["status"] = "error"
376+
flow["error"] = "Device flow timed out."
377+
378+
379+
async def _finalize_xai(
380+
flow_id: str, flow: dict, token_data: dict, client: httpx.AsyncClient
381+
):
382+
"""Persist xAI credentials and mark flow complete.
383+
384+
Credential shape mirrors XAiAuthBase._build_credentials_from_token_data
385+
so the existing provider loader picks them up unchanged.
386+
"""
387+
from rotator_library.providers.x_ai_auth_base import XAI_USERINFO_URL
388+
389+
access_token = token_data.get("access_token", "")
390+
refresh_token = token_data.get("refresh_token", "")
391+
392+
# Email discovery: id_token JWT → userinfo → sub fallback
393+
id_claims = _decode_jwt_payload(token_data.get("id_token", "")) or {}
394+
email = id_claims.get("email", "")
395+
sub = id_claims.get("sub", "")
396+
397+
if not email and access_token:
398+
try:
399+
userinfo_resp = await client.get(
400+
XAI_USERINFO_URL,
401+
headers={"Authorization": f"Bearer {access_token}"},
402+
timeout=10.0,
403+
)
404+
if userinfo_resp.is_success:
405+
userinfo = userinfo_resp.json()
406+
email = userinfo.get("email", "") or email
407+
sub = sub or userinfo.get("sub", "")
408+
except Exception as e:
409+
lib_logger.debug(f"xAI userinfo fetch failed: {e}")
410+
411+
if not email:
412+
email = sub or f"xai-user-{int(time.time())}"
413+
414+
expires_in = token_data.get("expires_in", 3600)
415+
416+
new_creds: Dict[str, Any] = {
417+
"access_token": access_token,
418+
"refresh_token": refresh_token,
419+
"expiry_date": time.time() + expires_in,
420+
"account_id": sub or email,
421+
"_proxy_metadata": {
422+
"email": email,
423+
"account_id": sub or email,
424+
"last_check_timestamp": time.time(),
425+
},
426+
}
427+
428+
_save_credential_file(flow, new_creds)
429+
flow["status"] = "complete"
430+
flow["result"] = {"login": email, "provider": "x-ai"}
431+
432+
270433
# ---------------------------------------------------------------------------
271434
# Codex / Gemini CLI / Anthropic: PKCE auth code + paste redirect URL
272435
# ---------------------------------------------------------------------------

0 commit comments

Comments
 (0)