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
+78 -5
View File
@@ -32,10 +32,12 @@ For captures / actions with `capture_after=True`:
from __future__ import annotations
import base64
import json
import logging
import os
import re
import struct
import sys
import threading
from typing import Any, Dict, List, Optional, Tuple
@@ -429,6 +431,61 @@ _DEFAULT_MAX_ELEMENTS = 100
# call passing a very large integer would silently disable the safeguard and
# reintroduce the original unbounded behavior.
_MAX_ALLOWED_MAX_ELEMENTS = 1000
_MIN_PROVIDER_IMAGE_DIMENSION = 8
def _image_dimensions_from_b64(image_b64: str) -> Optional[Tuple[int, int]]:
"""Return (width, height) for common inline screenshot formats.
Some providers reject images below 8x8 before the model sees the tool
result. Inspecting the encoded bytes here lets computer_use fall back to
its AX/SOM text payload instead of sending an unusable placeholder.
"""
if not image_b64:
return None
try:
raw = base64.b64decode(image_b64, validate=False)
except Exception:
return None
# PNG: signature + IHDR width/height.
if raw.startswith(b"\x89PNG\r\n\x1a\n") and len(raw) >= 24:
try:
width, height = struct.unpack(">II", raw[16:24])
return int(width), int(height)
except Exception:
return None
# JPEG: scan for SOF markers that carry dimensions.
if raw.startswith(b"\xff\xd8") and len(raw) > 4:
i = 2
while i + 9 < len(raw):
if raw[i] != 0xFF:
i += 1
continue
marker = raw[i + 1]
i += 2
while marker == 0xFF and i < len(raw):
marker = raw[i]
i += 1
if marker in {0xD8, 0xD9}:
continue
if marker == 0xDA:
break
if i + 2 > len(raw):
break
segment_len = int.from_bytes(raw[i:i + 2], "big")
if segment_len < 2 or i + segment_len > len(raw):
break
if marker in {
0xC0, 0xC1, 0xC2, 0xC3, 0xC5, 0xC6, 0xC7,
0xC9, 0xCA, 0xCB, 0xCD, 0xCE, 0xCF,
} and segment_len >= 7:
height = int.from_bytes(raw[i + 3:i + 5], "big")
width = int.from_bytes(raw[i + 5:i + 7], "big")
return int(width), int(height)
i += segment_len
return None
def _coerce_max_elements(value: Any) -> int:
@@ -457,6 +514,16 @@ def _capture_response(cap: CaptureResult, max_elements: int = _DEFAULT_MAX_ELEME
total_elements = len(cap.elements)
visible_elements = cap.elements[:max_elements]
truncated_elements = max(0, total_elements - len(visible_elements))
image_dimensions = _image_dimensions_from_b64(cap.png_b64 or "") if cap.png_b64 else None
response_width = image_dimensions[0] if image_dimensions else cap.width
response_height = image_dimensions[1] if image_dimensions else cap.height
image_too_small = bool(
image_dimensions
and (
image_dimensions[0] < _MIN_PROVIDER_IMAGE_DIMENSION
or image_dimensions[1] < _MIN_PROVIDER_IMAGE_DIMENSION
)
)
# Index only what's actually surfaced in the response — otherwise the
# human-readable summary references element indices the model cannot
@@ -464,7 +531,7 @@ def _capture_response(cap: CaptureResult, max_elements: int = _DEFAULT_MAX_ELEME
# 40-line index window).
element_index = _format_elements(visible_elements)
summary_lines = [
f"capture mode={cap.mode} {cap.width}x{cap.height}"
f"capture mode={cap.mode} {response_width}x{response_height}"
+ (f" app={cap.app}" if cap.app else "")
+ (f" window={cap.window_title!r}" if cap.window_title else ""),
f"{total_elements} interactable element(s):",
@@ -476,9 +543,15 @@ def _capture_response(cap: CaptureResult, max_elements: int = _DEFAULT_MAX_ELEME
# selected) has a valid value to hand to _route_capture_through_aux_vision.
# The AX path appends the "truncated to N of M" note to summary_lines
# below and rebuilds; the multimodal path keeps this version untouched.
if image_too_small:
summary_lines.append(
f" (screenshot omitted: {image_dimensions[0]}x{image_dimensions[1]} "
f"is below the {_MIN_PROVIDER_IMAGE_DIMENSION}x{_MIN_PROVIDER_IMAGE_DIMENSION} "
"provider minimum)"
)
summary = "\n".join(summary_lines)
if cap.png_b64 and cap.mode != "ax":
if cap.png_b64 and cap.mode != "ax" and not image_too_small:
# Decide whether to hand the screenshot to the auxiliary.vision
# pipeline (text-only result) or keep the multimodal envelope (main
# model handles vision natively). Issue #24015: previously the
@@ -510,7 +583,7 @@ def _capture_response(cap: CaptureResult, max_elements: int = _DEFAULT_MAX_ELEME
"image_url": {"url": f"data:{_mime};base64,{cap.png_b64}"}},
],
"text_summary": summary,
"meta": {"mode": cap.mode, "width": cap.width, "height": cap.height,
"meta": {"mode": cap.mode, "width": response_width, "height": response_height,
"elements": total_elements, "png_bytes": cap.png_bytes_len},
}
# AX-only (or image-missing fallback): text path actually carries the
@@ -523,8 +596,8 @@ def _capture_response(cap: CaptureResult, max_elements: int = _DEFAULT_MAX_ELEME
summary = "\n".join(summary_lines)
payload: Dict[str, Any] = {
"mode": cap.mode,
"width": cap.width,
"height": cap.height,
"width": response_width,
"height": response_height,
"app": cap.app,
"window_title": cap.window_title,
"elements": [_element_to_dict(e) for e in visible_elements],