fix(xai): accept Grok Build code during loopback wait + tiny screenshot guard

xAI's consent page renders the authorization code in-page instead of
redirecting to the loopback callback, so the listener just hangs and the
manual-paste flow demands a callback URL that never contains the token.

- auth.py: poll stdin non-blockingly while waiting for the xAI loopback
  callback; accept a pasted bare Grok Build code and substitute the locally
  generated state (PKCE code_verifier still binds the exchange). No need to
  wait for timeout or re-run with --manual-paste.
- computer_use: parse PNG/JPEG dimensions from base64 and fall back to the
  text/AX/SOM payload when the screenshot is below the provider minimum
  (8x8), which xAI rejects with HTTP 400.
- model_setup_flows.py: xAI credential reuse prompt uses the standard radio
  picker via a shared _prompt_auth_credentials_choice helper.
- main.py: thread a title through _prompt_provider_choice; re-home the helper
  import (flows live in model_setup_flows.py post-decomposition).

Salvaged from #36781 onto current main (contributor's main.py edits re-homed
to model_setup_flows.py, where the flows were extracted since the PR opened).
This commit is contained in:
Ondrej Drapalik
2026-06-09 23:21:24 -07:00
committed by Teknium
parent 095f526b11
commit 1c055a4c58
8 changed files with 330 additions and 50 deletions
+35 -3
View File
@@ -2665,12 +2665,23 @@ def _xai_wait_for_callback(
result: dict[str, Any],
*,
timeout_seconds: float = 180.0,
manual_paste_redirect_uri: Optional[str] = None,
) -> dict[str, Any]:
deadline = time.monotonic() + max(5.0, timeout_seconds)
if manual_paste_redirect_uri and sys.stdin.isatty():
print()
print("If xAI shows a Grok Build code instead of redirecting,")
print("paste that code here and press Enter.")
try:
while time.monotonic() < deadline:
if result["code"] or result["error"]:
return result
if manual_paste_redirect_uri:
raw_paste = _read_ready_stdin_line()
if raw_paste and raw_paste.strip():
pasted = _parse_pasted_callback(raw_paste)
pasted["_manual_paste"] = True
return pasted
time.sleep(0.1)
finally:
server.shutdown()
@@ -2694,6 +2705,21 @@ def _xai_wait_for_callback(
)
def _read_ready_stdin_line() -> Optional[str]:
"""Return one pending stdin line without blocking, if the terminal has one."""
try:
if not sys.stdin.isatty():
return None
import select
ready, _, _ = select.select([sys.stdin], [], [], 0)
if not ready:
return None
return sys.stdin.readline()
except Exception:
return None
def _spotify_token_payload_to_state(
token_payload: Dict[str, Any],
*,
@@ -6669,6 +6695,7 @@ def _xai_oauth_loopback_login(
authorization_endpoint = discovery["authorization_endpoint"]
token_endpoint = discovery["token_endpoint"]
allow_missing_state = False
if manual_paste:
# No HTTP listener — synthesize a redirect_uri matching what
# the server would have bound to so the authorize URL the user
@@ -6695,6 +6722,7 @@ def _xai_oauth_loopback_login(
print("Open this URL to authorize Hermes with xAI:")
print(authorize_url)
callback = _prompt_manual_callback_paste(redirect_uri)
allow_missing_state = True
else:
server, thread, callback_result, redirect_uri = _xai_start_callback_server()
try:
@@ -6734,6 +6762,7 @@ def _xai_oauth_loopback_login(
thread,
callback_result,
timeout_seconds=max(30.0, timeout_seconds * 9),
manual_paste_redirect_uri=redirect_uri,
)
except AuthError as exc:
if (
@@ -6750,6 +6779,7 @@ def _xai_oauth_loopback_login(
callback = _prompt_manual_callback_paste(redirect_uri)
if callback.get("code") is None and callback.get("error") is None:
raise exc
allow_missing_state = True
except Exception:
try:
server.shutdown()
@@ -6770,7 +6800,7 @@ def _xai_oauth_loopback_login(
code="xai_authorization_failed",
)
callback_state = callback.get("state")
# Manual-paste bare-code path: when a user pastes only the opaque
# Manual bare-code paths: when a user pastes only the opaque
# authorization code (no ``code=``/``state=`` query parameters),
# ``_parse_pasted_callback`` returns ``state=None``. xAI's consent
# page renders the code in-page rather than redirecting through the
@@ -6778,10 +6808,12 @@ def _xai_oauth_loopback_login(
# VPS, container consoles) the bare code is the only thing the user
# can obtain. PKCE (code_verifier) still binds the exchange to this
# client, so the local state-equality check is redundant on the
# bare-code path — we substitute the locally generated state to keep
# bare-code paths — we substitute the locally generated state to keep
# the rest of the validation chain (and the token exchange) unchanged.
# See #26923 (AccursedGalaxy comment, 2026-05-20).
if callback_state is None and manual_paste:
if callback.get("_manual_paste"):
allow_missing_state = True
if callback_state is None and (manual_paste or allow_missing_state):
callback_state = state
if callback_state != state:
raise AuthError(