diff --git a/README.md b/README.md
index c4204f01..5f84e4e2 100644
--- a/README.md
+++ b/README.md
@@ -49,7 +49,7 @@
- 🤖 Control Android and iOS devices with natural language commands
-- 🔀 Use OpenAI, Anthropic, Gemini, Ollama, DeepSeek, OpenRouter, and OpenAI-compatible models
+- 🔀 Use OpenAI, Anthropic, Gemini, xAI, Ollama, DeepSeek, OpenRouter, and OpenAI-compatible models
- 🧠 Run direct tasks or enable reasoning mode for complex multi-step automation
- 💻 Automate from the CLI, a terminal UI, Docker, or Python code
- 🐍 Extend agents with custom tools, structured output, app cards, and credentials
@@ -113,7 +113,7 @@ You should see confirmation that the Portal is installed and accessible.
mobilerun configure
```
-The wizard walks you through choosing a provider, auth method, and model. You can also use provider environment variables such as `GOOGLE_API_KEY`, `OPENAI_API_KEY`, `ANTHROPIC_API_KEY`, or `MINIMAX_API_KEY`.
+The wizard walks you through choosing a provider, auth method, and model. You can also use provider environment variables such as `GOOGLE_API_KEY`, `OPENAI_API_KEY`, `ANTHROPIC_API_KEY`, `XAI_API_KEY`, or `MINIMAX_API_KEY`.
### 4. Run your first command
diff --git a/docs/guides/cli.mdx b/docs/guides/cli.mdx
index 04a019a7..27f96e08 100644
--- a/docs/guides/cli.mdx
+++ b/docs/guides/cli.mdx
@@ -102,6 +102,12 @@ mobilerun run "Create shopping list" \
--provider OpenAI \
--model gpt-4o
+# xAI Grok
+export XAI_API_KEY=your-key
+mobilerun run "Open Settings" \
+ --provider XAI \
+ --model grok-4.5
+
# Anthropic Claude
export ANTHROPIC_API_KEY=your-key
mobilerun run "Reply to latest email" \
@@ -154,6 +160,7 @@ mobilerun run "Enable 2FA" \
|----------|---------|---------------------|
| GoogleGenAI | Included by default | `GOOGLE_API_KEY` |
| OpenAI | Included by default | `OPENAI_API_KEY` |
+| XAI | Included by default | `XAI_API_KEY` |
| OpenAILike | Included by default | Varies by provider |
| OpenRouter | Included by default | `OPENROUTER_API_KEY` |
| Ollama | Included by default | None (local) |
@@ -161,6 +168,27 @@ mobilerun run "Enable 2FA" \
| DeepSeek | Included by default | `DEEPSEEK_API_KEY` |
| MiniMax | Included by default | `MINIMAX_API_KEY` |
+### xAI API key and OAuth
+
+Configure XAI with an API key or OAuth:
+
+```bash
+# API key
+mobilerun configure \
+ --provider XAI \
+ --auth-mode api_key \
+ --model grok-4.5
+
+# OAuth through provider options
+mobilerun configure --provider XAI --auth-mode oauth --model grok-4.5
+
+# OAuth shortcut
+mobilerun configure xai
+
+# Device-code login for SSH/headless hosts
+mobilerun configure xai --device-code --no-browser
+```
+
### MiniMax endpoints and credentials
Run `mobilerun configure`, choose MiniMax, and then select the API region where
@@ -551,6 +579,7 @@ mobilerun ping --tcp
| `GOOGLE_API_KEY` | Google Gemini API key | None |
| `OPENAI_API_KEY` | OpenAI API key | None |
| `ANTHROPIC_API_KEY` | Anthropic API key | None |
+| `XAI_API_KEY` | xAI API key for Grok | None |
| `DEEPSEEK_API_KEY` | DeepSeek API key | None |
| `MINIMAX_API_KEY` | MiniMax API key | None |
| `MOBILERUN_CLOUD_API_KEY` | Mobilerun Cloud API key for `devices --cloud` and cloud device actions | None |
diff --git a/docs/quickstart.mdx b/docs/quickstart.mdx
index a543759c..5cabe91b 100644
--- a/docs/quickstart.mdx
+++ b/docs/quickstart.mdx
@@ -104,6 +104,9 @@ export OPENAI_API_KEY=your-api-key-here
# For Anthropic Claude
export ANTHROPIC_API_KEY=your-api-key-here
+
+# For xAI
+export XAI_API_KEY=your-api-key-here
```
### Run Your First Command via CLI
@@ -125,7 +128,7 @@ mobilerun run "Find a contact named John and send him an email" --reasoning
```
**Common CLI flags:**
-- `--provider` - LLM provider (GoogleGenAI, OpenAI, Anthropic, etc.)
+- `--provider` - LLM provider (GoogleGenAI, OpenAI, XAI, Anthropic, etc.)
- `--model` - Model name (gemini-3.5-flash-lite, gpt-5.5, etc.)
- `--vision` - Enable screenshot processing
- `--reasoning` - Enable multi-agent planning mode
diff --git a/docs/sdk/configuration.mdx b/docs/sdk/configuration.mdx
index b1571311..79f6bed6 100644
--- a/docs/sdk/configuration.mdx
+++ b/docs/sdk/configuration.mdx
@@ -611,7 +611,7 @@ mobilerun run "Task" --config /path/to/config.yaml
- `--config PATH` - Custom config file
- `--device SERIAL` - Device serial/IP
- `--agent NAME` - External agent to use. Not yet supported — reserved for future use.
-- `--provider PROVIDER` - LLM provider (OpenAI, Ollama, Anthropic, GoogleGenAI, DeepSeek, MiniMax)
+- `--provider PROVIDER` - LLM provider (OpenAI, XAI, Ollama, Anthropic, GoogleGenAI, DeepSeek, MiniMax)
- `--model MODEL` - LLM model name
- `--temperature FLOAT` - LLM temperature
- `--steps INT` - Max steps
@@ -635,6 +635,7 @@ Set API keys via environment variables:
export GOOGLE_API_KEY=your-key
export OPENAI_API_KEY=your-key
export ANTHROPIC_API_KEY=your-key
+export XAI_API_KEY=your-key
export DEEPSEEK_API_KEY=your-key
export MINIMAX_API_KEY=your-key
export MOBILERUN_CONFIG=/path/to/config.yaml # Custom config path
diff --git a/mobilerun/agent/providers/grok.py b/mobilerun/agent/providers/grok.py
new file mode 100644
index 00000000..d6c77bc9
--- /dev/null
+++ b/mobilerun/agent/providers/grok.py
@@ -0,0 +1,69 @@
+"""Shared Grok/xAI model and transport metadata."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, MutableMapping
+from typing import Any
+
+GROK_DEFAULT_MODEL = "grok-4.5"
+GROK_MODELS = (GROK_DEFAULT_MODEL,)
+GROK_MODEL_ALIASES = {
+ "grok-4.5-latest": GROK_DEFAULT_MODEL,
+ "grok-build-latest": GROK_DEFAULT_MODEL,
+}
+
+XAI_API_BASE = "https://api.x.ai/v1"
+GROK_CONTEXT_WINDOW = 500_000
+
+# Grok accepts temperature and top_p on the Responses API, but these legacy
+# Chat Completions controls are rejected. Filter after all constructor and
+# per-call kwargs are merged so an override cannot accidentally restore them.
+GROK_UNSUPPORTED_SAMPLING_PARAMS = frozenset(
+ {"presence_penalty", "frequency_penalty", "stop"}
+)
+
+
+def normalize_grok_model_id(model: object) -> str:
+ """Normalize public xAI/Grok aliases to Mobilerun's canonical model id."""
+
+ model_id = str(model or "").strip()
+ if model_id.startswith("xai/"):
+ model_id = model_id.removeprefix("xai/")
+ return GROK_MODEL_ALIASES.get(model_id, model_id)
+
+
+def sanitize_grok_responses_kwargs(
+ payload: MutableMapping[str, Any],
+ *,
+ omit_sampler_fields: bool = False,
+ omit_tool_choice: bool = False,
+) -> MutableMapping[str, Any]:
+ """Apply Grok's Responses API parameter contract to a final payload."""
+
+ filtered_params = set(GROK_UNSUPPORTED_SAMPLING_PARAMS)
+ if omit_sampler_fields:
+ filtered_params.update(("temperature", "top_p"))
+ if omit_tool_choice:
+ filtered_params.add("tool_choice")
+
+ for param in filtered_params:
+ payload.pop(param, None)
+ payload["store"] = False
+ payload.pop("reasoning", None)
+
+ # The OpenAI SDK merges ``extra_body`` after its normal typed parameters,
+ # so an unsanitized value here could otherwise restore storage, reasoning,
+ # an unsupported sampler, or even a caller-selected model. Preserve other
+ # extension fields while removing every value Mobilerun pins or rejects.
+ extra_body = payload.get("extra_body")
+ if extra_body is not None:
+ if not isinstance(extra_body, Mapping):
+ payload.pop("extra_body", None)
+ else:
+ sanitized_extra_body = dict(extra_body)
+ for param in filtered_params:
+ sanitized_extra_body.pop(param, None)
+ for param in ("model", "store", "reasoning"):
+ sanitized_extra_body.pop(param, None)
+ payload["extra_body"] = sanitized_extra_body
+ return payload
diff --git a/mobilerun/agent/providers/registry.py b/mobilerun/agent/providers/registry.py
index 54c77e11..7d9d69f4 100644
--- a/mobilerun/agent/providers/registry.py
+++ b/mobilerun/agent/providers/registry.py
@@ -6,6 +6,12 @@
ANTHROPIC_OAUTH_DEFAULT_MODEL,
ANTHROPIC_OAUTH_MODELS,
)
+from mobilerun.agent.providers.grok import (
+ GROK_DEFAULT_MODEL,
+ GROK_MODELS,
+ XAI_API_BASE,
+ normalize_grok_model_id,
+)
from mobilerun.agent.providers.minimax import MINIMAX_GLOBAL_BASE_URL
from mobilerun.agent.providers.types import (
ProviderFamilySpec,
@@ -14,6 +20,7 @@
from mobilerun.config_manager.credential_paths import (
ANTHROPIC_OAUTH_CREDENTIAL_PATH,
GEMINI_OAUTH_CREDENTIAL_PATH,
+ GROK_OAUTH_CREDENTIAL_PATH,
OPENAI_OAUTH_CREDENTIAL_PATH,
)
@@ -22,6 +29,7 @@
VARIANT_ENV_KEY_SLOT: dict[str, str] = {
"GoogleGenAI": "google",
"OpenAIResponses": "openai",
+ "XAI": "xai",
"Anthropic": "anthropic",
"ZAI": "zai",
"ZAI_Coding": "zai",
@@ -133,6 +141,29 @@
),
),
),
+ ProviderFamilySpec(
+ id="xai",
+ display_name="XAI",
+ variants=(
+ ProviderVariantSpec(
+ id="XAI",
+ runtime_provider_name="XAI",
+ auth_mode="api_key",
+ default_model=GROK_DEFAULT_MODEL,
+ models=GROK_MODELS,
+ requires_api_key=True,
+ base_url=XAI_API_BASE,
+ ),
+ ProviderVariantSpec(
+ id="xai_oauth",
+ runtime_provider_name="xai_oauth",
+ auth_mode="oauth",
+ default_model=GROK_DEFAULT_MODEL,
+ models=GROK_MODELS,
+ credential_path=str(GROK_OAUTH_CREDENTIAL_PATH),
+ ),
+ ),
+ ),
ProviderFamilySpec(
id="ollama",
display_name="Ollama",
@@ -285,6 +316,8 @@ def normalize_model_id_for_variant(
if family_id == "openai":
candidate = OPENAI_MODEL_ALIASES.get(candidate, candidate)
+ elif family_id == "xai":
+ candidate = normalize_grok_model_id(candidate)
if candidate in allowed_model_ids:
return candidate
diff --git a/mobilerun/agent/providers/setup_service.py b/mobilerun/agent/providers/setup_service.py
index bbe172d0..25a3a1de 100644
--- a/mobilerun/agent/providers/setup_service.py
+++ b/mobilerun/agent/providers/setup_service.py
@@ -23,6 +23,7 @@
# which allocates the full KV cache (256K-context models -> ~19 GB) and
# spills to CPU on typical machines. -1 restores model-max for big GPUs.
"Ollama": {"context_window": 32768},
+ "XAI": {"context_window": 500_000},
}
HIDDEN_ROLE_FALLBACKS: tuple[str, ...] = ("app_opener", "structured_output")
@@ -176,10 +177,12 @@ def create_profile_for_variant(
api_key_source=selection.api_key_source,
base_url=base_url,
api_base=(
- base_url if runtime_provider_name in {"OpenAILike", "MiniMax"} else None
+ base_url
+ if runtime_provider_name in {"OpenAILike", "MiniMax", "XAI"}
+ else None
),
credential_path=selection.credential_path or variant.credential_path,
- kwargs=kwargs if env_slot is None else {},
+ kwargs=kwargs if env_slot is None or variant.id == "XAI" else {},
)
diff --git a/mobilerun/agent/usage.py b/mobilerun/agent/usage.py
index 13714a3b..36a5c39d 100644
--- a/mobilerun/agent/usage.py
+++ b/mobilerun/agent/usage.py
@@ -28,6 +28,8 @@
"MobilerunAnthropic": "Anthropic",
"MobilerunOpenAIResponses": "OpenAIResponses",
"openai_responses_llm": "OpenAIResponses",
+ "GrokOAuth": "OpenAIResponses",
+ "xai_oauth": "OpenAIResponses",
"Ollama_llm": "Ollama",
}
@@ -57,6 +59,14 @@ def _usage_field(usage: Any, *names: str) -> int:
return 0
+def _response_field(response: Any, name: str) -> Any:
+ """Read a response field from either an SDK object or decoded JSON."""
+
+ if isinstance(response, dict):
+ return response.get(name)
+ return getattr(response, name, None)
+
+
def _normalize_provider_name(provider: str) -> str:
return PROVIDER_ALIASES.get(provider, provider)
@@ -64,7 +74,7 @@ def _normalize_provider_name(provider: str) -> str:
def get_usage_from_response(provider: str, chat_rsp: ChatResponse) -> UsageResult:
provider = _normalize_provider_name(provider)
rsp = chat_rsp.raw
- if not rsp:
+ if not rsp and provider not in ("OpenAIResponses", "OpenAIOAuth"):
raise ValueError("No raw response in chat response")
if provider in {
@@ -100,15 +110,25 @@ def get_usage_from_response(provider: str, chat_rsp: ChatResponse) -> UsageResul
requests=1,
)
elif provider in ("OpenAIResponses", "OpenAIOAuth"):
- usage = getattr(rsp, "usage", None)
+ usage = _response_field(rsp, "usage")
+ if usage is None:
+ # Streaming Responses end with a ``response.completed`` event. Its
+ # accounting belongs to the nested final Response rather than the
+ # event itself.
+ usage = _response_field(_response_field(rsp, "response"), "usage")
+ if usage is None:
+ # LlamaIndex also copies completed-stream usage into the final
+ # ChatResponse's additional kwargs. Keep this fallback for custom
+ # or normalized event representations that omit the raw response.
+ usage = _response_field(chat_rsp.additional_kwargs, "usage")
if usage is None:
return UsageResult(
request_tokens=0, response_tokens=0, total_tokens=0, requests=1
)
return UsageResult(
- request_tokens=getattr(usage, "input_tokens", 0) or 0,
- response_tokens=getattr(usage, "output_tokens", 0) or 0,
- total_tokens=getattr(usage, "total_tokens", 0) or 0,
+ request_tokens=_usage_field(usage, "input_tokens"),
+ response_tokens=_usage_field(usage, "output_tokens"),
+ total_tokens=_usage_field(usage, "total_tokens"),
requests=1,
)
elif provider in {"Anthropic", "Anthropic_LLM", "AnthropicOAuthLLM"}:
diff --git a/mobilerun/agent/utils/llm_picker.py b/mobilerun/agent/utils/llm_picker.py
index 59a3df2e..2aa92fe3 100644
--- a/mobilerun/agent/utils/llm_picker.py
+++ b/mobilerun/agent/utils/llm_picker.py
@@ -8,6 +8,13 @@
anthropic_model_context_window,
anthropic_model_omits_sampling_params,
)
+from mobilerun.agent.providers.grok import (
+ GROK_CONTEXT_WINDOW,
+ GROK_DEFAULT_MODEL,
+ XAI_API_BASE,
+ normalize_grok_model_id,
+ sanitize_grok_responses_kwargs,
+)
from mobilerun.agent.providers.minimax import (
MINIMAX_GLOBAL_BASE_URL,
warn_if_legacy_minimax_endpoint,
@@ -34,6 +41,7 @@
"DeepSeek",
"OpenRouter",
"MiniMax",
+ "XAI",
]
PROVIDER_ALIASES = {
@@ -49,6 +57,7 @@
"openai_like": "OpenAILike",
"zai": "ZAI",
"z.ai": "ZAI",
+ "xai": "XAI",
}
ZAI_GLOBAL_API_BASE = "https://api.z.ai/api/paas/v4"
@@ -184,12 +193,29 @@ def _prepare_ollama_kwargs(kwargs: dict[str, Any], llm_class: Any) -> dict[str,
return kwargs
-def _load_openai_responses(**kwargs: Any) -> LLM:
+def _load_openai_responses(*, grok: bool = False, **kwargs: Any) -> LLM:
from llama_index.llms.openai.responses import OpenAIResponses
+ from llama_index.llms.openai.utils import to_openai_message_dicts
class MobilerunOpenAIResponses(OpenAIResponses):
- def _sanitize_call_kwargs(self, call_kwargs: dict[str, Any]) -> dict[str, Any]:
+ def _sanitize_call_kwargs(
+ self,
+ call_kwargs: dict[str, Any],
+ *,
+ omit_tool_choice: bool = False,
+ ) -> dict[str, Any]:
sanitized = dict(call_kwargs)
+ if grok:
+ sanitized = dict(
+ sanitize_grok_responses_kwargs(
+ sanitized,
+ omit_tool_choice=omit_tool_choice,
+ )
+ )
+ # Runtime and additional kwargs are merged after constructor
+ # defaults. Re-pin the canonical model after that final merge.
+ sanitized["model"] = self.model
+ return sanitized
effective_model = sanitized.get("model", self.model)
if _openai_responses_model_omits_sampling_params(effective_model):
for param in OPENAI_RESPONSES_UNSUPPORTED_SAMPLING_PARAMS:
@@ -199,6 +225,21 @@ def _sanitize_call_kwargs(self, call_kwargs: dict[str, Any]) -> dict[str, Any]:
def _get_model_kwargs(self, **kwargs: Any) -> dict[str, Any]:
return self._sanitize_call_kwargs(super()._get_model_kwargs(**kwargs))
+ def _sanitize_structured_call_kwargs(
+ self, call_kwargs: dict[str, Any]
+ ) -> dict[str, Any]:
+ sanitized = self._sanitize_call_kwargs(call_kwargs, omit_tool_choice=grok)
+ if grok:
+ # The upstream structured adapter passes ``store=self.store``
+ # explicitly. The constructor already pins that field false.
+ # xAI also rejects tool_choice when no tools are supplied, so
+ # structured parsing must not allow either the upstream default
+ # or a generic caller override to restore it.
+ sanitized.pop("store", None)
+ sanitized.pop("model", None)
+ sanitized.pop("tool_choice", None)
+ return sanitized
+
def structured_predict(
self,
output_cls: Any,
@@ -206,10 +247,29 @@ def structured_predict(
llm_kwargs: dict[str, Any] | None = None,
**prompt_args: Any,
) -> Any:
+ if grok:
+ messages = prompt.format_messages(**prompt_args)
+ message_dicts = to_openai_message_dicts(
+ messages, model=self.model, is_responses_api=True
+ )
+ response = self._client.responses.parse(
+ model=self._responses_model,
+ input=message_dicts,
+ text_format=output_cls,
+ store=self.store,
+ **self._sanitize_structured_call_kwargs(dict(llm_kwargs or {})),
+ )
+ if response.output_parsed is not None:
+ return response.output_parsed
+ raise ValueError(
+ "Failed to produce a structured response from the model."
+ )
return super().structured_predict(
output_cls,
prompt,
- llm_kwargs=self._sanitize_call_kwargs(dict(llm_kwargs or {})),
+ llm_kwargs=self._sanitize_structured_call_kwargs(
+ dict(llm_kwargs or {})
+ ),
**prompt_args,
)
@@ -220,13 +280,41 @@ async def astructured_predict(
llm_kwargs: dict[str, Any] | None = None,
**prompt_args: Any,
) -> Any:
+ if grok:
+ messages = prompt.format_messages(**prompt_args)
+ message_dicts = to_openai_message_dicts(
+ messages, model=self.model, is_responses_api=True
+ )
+ response = await self._aclient.responses.parse(
+ model=self._responses_model,
+ input=message_dicts,
+ text_format=output_cls,
+ store=self.store,
+ **self._sanitize_structured_call_kwargs(dict(llm_kwargs or {})),
+ )
+ if response.output_parsed is not None:
+ return response.output_parsed
+ raise ValueError(
+ "Failed to produce a structured response from the model."
+ )
return await super().astructured_predict(
output_cls,
prompt,
- llm_kwargs=self._sanitize_call_kwargs(dict(llm_kwargs or {})),
+ llm_kwargs=self._sanitize_structured_call_kwargs(
+ dict(llm_kwargs or {})
+ ),
**prompt_args,
)
+ if grok:
+ kwargs = dict(sanitize_grok_responses_kwargs(kwargs))
+ additional_kwargs = dict(kwargs.get("additional_kwargs") or {})
+ sanitize_grok_responses_kwargs(additional_kwargs)
+ # ``store`` is a first-class adapter field; keeping a second copy in
+ # additional_kwargs is unnecessary and makes intent harder to inspect.
+ additional_kwargs.pop("store", None)
+ kwargs["additional_kwargs"] = additional_kwargs
+
filtered_kwargs = {k: v for k, v in kwargs.items() if v is not None}
logger.debug(
"Initializing MobilerunOpenAIResponses with kwargs: "
@@ -431,6 +519,10 @@ def load_llm(provider_name: str, model: str | None = None, **kwargs: Any) -> LLM
if model is not None:
if provider_name == "OpenAIResponses":
model = normalize_model_id_for_variant("openai", "api_key", model)
+ elif provider_name == "XAI":
+ model = normalize_grok_model_id(model)
+ elif provider_name == "xai_oauth":
+ model = normalize_model_id_for_variant("xai", "oauth", model)
elif provider_name == "openai_oauth":
model = normalize_model_id_for_variant("openai", "oauth", model)
kwargs["model"] = model
@@ -460,6 +552,10 @@ def load_llm(provider_name: str, model: str | None = None, **kwargs: Any) -> LLM
return GeminiOAuthCodeAssistLLM(
**{k: v for k, v in kwargs.items() if v is not None}
)
+ if provider_name == "xai_oauth":
+ from mobilerun.agent.utils.oauth.grok_oauth_llm import GrokOAuth
+
+ return GrokOAuth(**{k: v for k, v in kwargs.items() if v is not None})
# Legacy aliases: MiniMax and DeepSeek route through OpenAILike.
if provider_name == "MiniMax":
@@ -489,6 +585,36 @@ def load_llm(provider_name: str, model: str | None = None, **kwargs: Any) -> LLM
kwargs["api_base"] = kwargs.pop("base_url")
kwargs.setdefault("api_base", ZAI_GLOBAL_API_BASE)
+ if provider_name == "XAI":
+ import os
+
+ # Mobilerun's reasoning mode selects its agent architecture. It does
+ # not opt Grok into a provider-specific reasoning effort.
+ kwargs.pop("reasoning_options", None)
+ api_key = kwargs.get("api_key")
+ if not isinstance(api_key, str) or not api_key.strip():
+ api_key = os.environ.get("XAI_API_KEY")
+ if not isinstance(api_key, str) or not api_key.strip():
+ raise ValueError(
+ "XAI requires an API key. Pass api_key explicitly or set "
+ "XAI_API_KEY."
+ )
+
+ kwargs["api_key"] = api_key
+ # The lowercase runtime alias should be useful without a separately
+ # generated profile. Keep its implicit model aligned with the catalog.
+ kwargs.setdefault("model", GROK_DEFAULT_MODEL)
+ # XAI_API_KEY must only ever be sent to xAI's pinned endpoint. Ignore
+ # generic CLI/profile URL overrides rather than allowing a malicious
+ # config to redirect the bearer credential to another host.
+ kwargs.pop("base_url", None)
+ kwargs["api_base"] = XAI_API_BASE
+ # Grok 4.5's catalog context is provider metadata, not a caller-tunable
+ # endpoint option. Keep hand-written runtime profiles aligned with the
+ # first-class provider catalog as well as generated profiles.
+ kwargs["context_window"] = GROK_CONTEXT_WINDOW
+ return _load_openai_responses(grok=True, **kwargs)
+
if provider_name == "DeepSeek":
import os
diff --git a/mobilerun/agent/utils/oauth/anthropic_oauth_llm.py b/mobilerun/agent/utils/oauth/anthropic_oauth_llm.py
index f371cd9d..92096543 100644
--- a/mobilerun/agent/utils/oauth/anthropic_oauth_llm.py
+++ b/mobilerun/agent/utils/oauth/anthropic_oauth_llm.py
@@ -8,7 +8,8 @@
import threading
import time
import webbrowser
-from http.server import BaseHTTPRequestHandler, HTTPServer
+from http.server import BaseHTTPRequestHandler
+from http.server import ThreadingHTTPServer as HTTPServer
from pathlib import Path
from typing import Any, Dict, Literal, Optional, Sequence
from urllib.parse import parse_qs, urlencode, urlparse
@@ -37,6 +38,11 @@
anthropic_model_omits_sampling_params,
strip_anthropic_sampling_params,
)
+from mobilerun.agent.utils.oauth.login_timeout import (
+ OAuthLoginDeadline,
+ open_browser_async,
+)
+from mobilerun.config_manager.auth_profile_store import AuthProfileStore
from mobilerun.config_manager.credential_paths import ANTHROPIC_OAUTH_CREDENTIAL_PATH
DEFAULT_MODEL = ANTHROPIC_OAUTH_DEFAULT_MODEL
@@ -262,39 +268,29 @@ def _load_credentials_from_file(self, credential_path: str) -> None:
if isinstance(expires_at, (int, float)):
self._access_token_expiry = float(expires_at) / 1000.0
- def _persist_credentials(self) -> None:
+ def _persist_credentials(
+ self,
+ *,
+ deadline: OAuthLoginDeadline | None = None,
+ ) -> None:
if not self.credential_path:
return
path = Path(self.credential_path).expanduser()
- path.parent.mkdir(parents=True, exist_ok=True)
-
- existing: Dict[str, Any] = {}
- if path.exists():
- try:
- loaded = json.loads(path.read_text(encoding="utf-8"))
- if isinstance(loaded, dict):
- existing = loaded
- except Exception:
- existing = {}
-
- existing["claudeAiOauth"] = {
- "accessToken": self._cached_access_token,
- "refreshToken": self._cached_refresh_token,
- "expiresAt": (
- int(self._access_token_expiry * 1000)
- if self._access_token_expiry
- else None
- ),
- "scopes": self.refresh_scope.split(),
- }
-
- tmp_path = path.with_suffix(path.suffix + ".tmp")
- tmp_path.write_text(json.dumps(existing, indent=2), encoding="utf-8")
- os.replace(tmp_path, path)
- try:
- os.chmod(path, 0o600)
- except OSError:
- pass
+ AuthProfileStore(path).update_slot(
+ "claudeAiOauth",
+ {
+ "accessToken": self._cached_access_token,
+ "refreshToken": self._cached_refresh_token,
+ "expiresAt": (
+ int(self._access_token_expiry * 1000)
+ if self._access_token_expiry
+ else None
+ ),
+ "scopes": self.refresh_scope.split(),
+ },
+ lock_timeout=deadline.remaining() if deadline is not None else None,
+ before_commit=deadline.check if deadline is not None else None,
+ )
def _token_headers(self) -> Dict[str, str]:
headers = {
@@ -358,6 +354,7 @@ def _exchange_authorization_code(
code_verifier: str,
state: str,
expires_in: Optional[int] = None,
+ deadline: OAuthLoginDeadline | None = None,
) -> str:
request_body: Dict[str, Any] = {
"grant_type": "authorization_code",
@@ -370,14 +367,21 @@ def _exchange_authorization_code(
if expires_in is not None:
request_body["expires_in"] = expires_in
+ request_timeout = (
+ self.timeout if deadline is None else deadline.remaining(cap=self.timeout)
+ )
res = self._session.post(
self.token_url,
headers=self._token_headers(),
json=request_body,
- timeout=self.timeout,
+ timeout=request_timeout,
)
+ if deadline is not None:
+ deadline.check()
res.raise_for_status()
data = res.json()
+ if deadline is not None:
+ deadline.check()
access_token = data.get("access_token")
if not isinstance(access_token, str) or not access_token:
@@ -386,8 +390,11 @@ def _exchange_authorization_code(
)
refresh_token = data.get("refresh_token")
- if isinstance(refresh_token, str) and refresh_token:
- self._cached_refresh_token = refresh_token
+ cached_refresh_token = (
+ refresh_token
+ if isinstance(refresh_token, str) and refresh_token
+ else self._cached_refresh_token
+ )
expires_in = data.get("expires_in", 28_800)
try:
@@ -395,9 +402,25 @@ def _exchange_authorization_code(
except (TypeError, ValueError):
expires_in_s = 28_800
+ if deadline is not None:
+ deadline.check()
+ previous = (
+ self._cached_access_token,
+ self._cached_refresh_token,
+ self._access_token_expiry,
+ )
+ self._cached_refresh_token = cached_refresh_token
self._cached_access_token = access_token
self._access_token_expiry = time.time() + expires_in_s
- self._persist_credentials()
+ try:
+ self._persist_credentials(deadline=deadline)
+ except BaseException:
+ (
+ self._cached_access_token,
+ self._cached_refresh_token,
+ self._access_token_expiry,
+ ) = previous
+ raise
return access_token
def _build_auth_url(
@@ -429,7 +452,11 @@ def login(
callback_port: int = 0,
callback_path: str = "/callback",
expires_in: Optional[int] = None,
+ deadline: OAuthLoginDeadline | None = None,
) -> str:
+ active_deadline = deadline or OAuthLoginDeadline(timeout_seconds)
+ active_deadline.check()
+
# Headless environments: skip local server, use hosted callback page
use_headless = _is_headless_environment() or os.environ.get(
"DROIDRUN_OAUTH_MANUAL", ""
@@ -439,11 +466,13 @@ def login(
open_browser=open_browser,
timeout_seconds=timeout_seconds,
expires_in=expires_in,
+ deadline=active_deadline,
)
# Desktop: browser callback server
result: Dict[str, Optional[str]] = {"code": None, "state": None, "error": None}
done = threading.Event()
+ callback_lock = threading.Lock()
code_verifier, code_challenge = _pkce_pair()
state = _b64_no_pad(secrets.token_bytes(32))
@@ -460,12 +489,17 @@ def do_GET(self) -> None: # noqa: N802
self.end_headers()
return
- params = parse_qs(parsed.query)
- result["code"] = params.get("code", [None])[0]
- result["state"] = params.get("state", [None])[0]
- result["error"] = params.get("error", [None])[0]
-
- ok = result["code"] is not None and result["error"] is None
+ with callback_lock:
+ if done.is_set():
+ self.send_response(409)
+ self.end_headers()
+ return
+ params = parse_qs(parsed.query)
+ result["code"] = params.get("code", [None])[0]
+ result["state"] = params.get("state", [None])[0]
+ result["error"] = params.get("error", [None])[0]
+ ok = result["code"] is not None and result["error"] is None
+ done.set()
self.send_response(200 if ok else 400)
self.send_header("Content-Type", "text/html; charset=utf-8")
self.end_headers()
@@ -477,7 +511,6 @@ def do_GET(self) -> None: # noqa: N802
self.wfile.write(
b"Login failed. Return to your terminal.
"
)
- done.set()
def log_message(self, format: str, *args: Any) -> None: # noqa: A003
return
@@ -494,8 +527,10 @@ def log_message(self, format: str, *args: Any) -> None: # noqa: A003
open_browser=open_browser,
timeout_seconds=timeout_seconds,
expires_in=expires_in,
+ deadline=active_deadline,
)
+ active_deadline.check()
actual_port = httpd.server_address[1]
redirect_uri = f"http://localhost:{actual_port}{callback_path}"
auth_url = self._build_auth_url(
@@ -505,17 +540,19 @@ def log_message(self, format: str, *args: Any) -> None: # noqa: A003
)
server_thread = threading.Thread(target=httpd.serve_forever, daemon=True)
+ httpd.daemon_threads = True
server_thread.start()
try:
print(f"Open this URL to login:\n{auth_url}\n")
if open_browser:
- webbrowser.open(auth_url)
+ open_browser_async(auth_url, webbrowser.open)
- if not done.wait(timeout=timeout_seconds):
+ if not done.wait(timeout=active_deadline.remaining()):
raise TimeoutError(
"OAuth login timed out before callback was received."
)
+ active_deadline.check()
if result["error"]:
raise RuntimeError(f"OAuth callback returned error: {result['error']}")
@@ -526,17 +563,19 @@ def log_message(self, format: str, *args: Any) -> None: # noqa: A003
"OAuth callback did not include an authorization code."
)
- return self._exchange_authorization_code(
+ access_token = self._exchange_authorization_code(
code=result["code"],
redirect_uri=redirect_uri,
code_verifier=code_verifier,
state=state,
expires_in=expires_in,
+ deadline=active_deadline,
)
finally:
self.authorize_url = original_authorize_url
httpd.shutdown()
httpd.server_close()
+ return access_token
def login_headless(
self,
@@ -545,8 +584,12 @@ def login_headless(
timeout_seconds: float = 300.0,
input_fn: Any = input,
expires_in: Optional[int] = None,
+ deadline: OAuthLoginDeadline | None = None,
) -> str:
"""Headless OAuth flow for SSH/WSL environments."""
+ active_deadline = deadline or OAuthLoginDeadline(timeout_seconds)
+ active_deadline.check()
+
code_verifier, code_challenge = _pkce_pair()
state = _b64_no_pad(secrets.token_bytes(32))
redirect_uri = "https://platform.claude.com/oauth/code/callback"
@@ -568,9 +611,8 @@ def login_headless(
f"\n2. Complete sign-in, then paste the authorization code shown on the page.\n"
)
if open_browser:
- webbrowser.open(auth_url)
+ open_browser_async(auth_url, webbrowser.open)
- deadline = time.time() + timeout_seconds
input_queue: queue.Queue[Optional[str]] = queue.Queue()
stop = threading.Event()
need_more = threading.Event()
@@ -591,16 +633,13 @@ def _reader() -> None:
try:
for attempt in range(_MAX_CODE_ATTEMPTS):
- remaining = deadline - time.time()
- if remaining <= 0:
- raise TimeoutError("OAuth login timed out.")
-
need_more.set()
try:
- raw = input_queue.get(timeout=remaining)
+ raw = input_queue.get(timeout=active_deadline.remaining())
except queue.Empty:
raise TimeoutError("OAuth login timed out.") from None
+ active_deadline.check()
if raw is None:
raise RuntimeError("Login failed — stdin closed.")
@@ -623,6 +662,7 @@ def _reader() -> None:
code_verifier=code_verifier,
state=state,
expires_in=expires_in,
+ deadline=active_deadline,
)
if attempt == 0:
print("Invalid code. Try again.")
diff --git a/mobilerun/agent/utils/oauth/gemini_oauth_code_assist_llm.py b/mobilerun/agent/utils/oauth/gemini_oauth_code_assist_llm.py
index 8be5fbe1..6b7b9749 100644
--- a/mobilerun/agent/utils/oauth/gemini_oauth_code_assist_llm.py
+++ b/mobilerun/agent/utils/oauth/gemini_oauth_code_assist_llm.py
@@ -8,7 +8,8 @@
import threading
import time
import webbrowser
-from http.server import BaseHTTPRequestHandler, HTTPServer
+from http.server import BaseHTTPRequestHandler
+from http.server import ThreadingHTTPServer as HTTPServer
from pathlib import Path
from typing import Any, ClassVar, Dict, Optional, Sequence
from urllib.parse import parse_qs, urlencode, urlparse
@@ -31,6 +32,11 @@
from llama_index.core.llms.callbacks import llm_chat_callback, llm_completion_callback
from llama_index.core.llms.custom import CustomLLM
+from mobilerun.agent.utils.oauth.login_timeout import (
+ OAuthLoginDeadline,
+ open_browser_async,
+)
+from mobilerun.config_manager.auth_profile_store import AuthProfileStore
from mobilerun.config_manager.credential_paths import GEMINI_OAUTH_CREDENTIAL_PATH
DEFAULT_MODEL = "gemini-3.5-flash-low"
@@ -291,40 +297,30 @@ def _load_credentials_from_file(self, credential_path: str) -> None:
if isinstance(expiry_ms, (int, float)):
self._access_token_expiry = float(expiry_ms) / 1000.0
- def _persist_credentials(self) -> None:
+ def _persist_credentials(
+ self,
+ *,
+ deadline: OAuthLoginDeadline | None = None,
+ ) -> None:
if not self.credential_path:
return
path = Path(self.credential_path).expanduser()
- path.parent.mkdir(parents=True, exist_ok=True)
-
- existing: Dict[str, Any] = {}
- if path.exists():
- try:
- loaded = json.loads(path.read_text(encoding="utf-8"))
- if isinstance(loaded, dict):
- existing = loaded
- except Exception:
- existing = {}
-
- existing[self.credential_slot] = {
- "access_token": self._cached_access_token,
- "refresh_token": self._cached_refresh_token,
- "token_type": "Bearer",
- "expiry_date": (
- int(self._access_token_expiry * 1000)
- if self._access_token_expiry
- else None
- ),
- }
-
- tmp_path = path.with_suffix(path.suffix + ".tmp")
- tmp_path.write_text(json.dumps(existing, indent=2), encoding="utf-8")
- os.replace(tmp_path, path)
- try:
- os.chmod(path, 0o600)
- except OSError:
- pass
+ AuthProfileStore(path).update_slot(
+ self.credential_slot,
+ {
+ "access_token": self._cached_access_token,
+ "refresh_token": self._cached_refresh_token,
+ "token_type": "Bearer",
+ "expiry_date": (
+ int(self._access_token_expiry * 1000)
+ if self._access_token_expiry
+ else None
+ ),
+ },
+ lock_timeout=deadline.remaining() if deadline is not None else None,
+ before_commit=deadline.check if deadline is not None else None,
+ )
def _metadata_payload(self) -> Dict[str, str]:
return {
@@ -343,7 +339,12 @@ def _build_headers(self, token: str) -> Dict[str, str]:
"Client-Metadata": json.dumps(self._metadata_payload()),
}
- def fetch_available_models(self) -> list[Dict[str, Any]]:
+ def fetch_available_models(
+ self,
+ *,
+ deadline: OAuthLoginDeadline | None = None,
+ access_token: str | None = None,
+ ) -> list[Dict[str, Any]]:
"""Agent-usable Gemini models for the current entitlement.
Calls Code Assist ``fetchAvailableModels`` and returns dicts with
@@ -351,15 +352,22 @@ def fetch_available_models(self) -> list[Dict[str, Any]]:
deprecated ids are filtered out. Used to verify login and to optionally
discover the live catalog. Requires the Antigravity client headers.
"""
- token = self._resolve_access_token()
+ token = access_token or self._resolve_access_token()
+ effective_timeout = (
+ self.timeout if deadline is None else deadline.remaining(cap=self.timeout)
+ )
response = self._session.post(
self._method_url(DEFAULT_CODE_ASSIST_MODELS_METHOD),
headers=self._build_headers(token),
json={},
- timeout=self.timeout,
+ timeout=effective_timeout,
)
+ if deadline is not None:
+ deadline.check()
response.raise_for_status()
data = response.json()
+ if deadline is not None:
+ deadline.check()
models = data.get("models")
if not isinstance(models, dict):
return []
@@ -379,6 +387,8 @@ def fetch_available_models(self) -> list[Dict[str, Any]]:
"supports_images": bool(meta.get("supportsImages")),
}
)
+ if deadline is not None:
+ deadline.check()
return out
def _access_token_is_stale(self) -> bool:
@@ -430,6 +440,8 @@ def _exchange_authorization_code(
code: str,
redirect_uri: str,
code_verifier: Optional[str] = None,
+ deadline: OAuthLoginDeadline | None = None,
+ persist_credentials: bool = True,
) -> str:
payload = {
"grant_type": "authorization_code",
@@ -441,13 +453,20 @@ def _exchange_authorization_code(
if code_verifier:
payload["code_verifier"] = code_verifier
+ request_timeout = (
+ self.timeout if deadline is None else deadline.remaining(cap=self.timeout)
+ )
response = self._session.post(
self.token_url,
data=payload,
- timeout=self.timeout,
+ timeout=request_timeout,
)
+ if deadline is not None:
+ deadline.check()
response.raise_for_status()
data = response.json()
+ if deadline is not None:
+ deadline.check()
access_token = data.get("access_token")
if not isinstance(access_token, str) or not access_token:
@@ -456,8 +475,11 @@ def _exchange_authorization_code(
)
refresh_token = data.get("refresh_token")
- if isinstance(refresh_token, str) and refresh_token:
- self._cached_refresh_token = refresh_token
+ cached_refresh_token = (
+ refresh_token
+ if isinstance(refresh_token, str) and refresh_token
+ else self._cached_refresh_token
+ )
expires_in = data.get("expires_in", 3600)
try:
@@ -465,9 +487,26 @@ def _exchange_authorization_code(
except (TypeError, ValueError):
expires_in_s = 3600
+ if deadline is not None:
+ deadline.check()
+ previous = (
+ self._cached_access_token,
+ self._cached_refresh_token,
+ self._access_token_expiry,
+ )
+ self._cached_refresh_token = cached_refresh_token
self._cached_access_token = access_token
self._access_token_expiry = time.time() + expires_in_s
- self._persist_credentials()
+ try:
+ if persist_credentials:
+ self._persist_credentials(deadline=deadline)
+ except BaseException:
+ (
+ self._cached_access_token,
+ self._cached_refresh_token,
+ self._access_token_expiry,
+ ) = previous
+ raise
return access_token
def _build_auth_url(
@@ -502,7 +541,12 @@ def login(
callback_port: int = 0,
callback_path: str = "/oauth2callback",
prompt_consent: bool = True,
+ deadline: OAuthLoginDeadline | None = None,
+ persist_credentials: bool = True,
) -> str:
+ active_deadline = deadline or OAuthLoginDeadline(timeout_seconds)
+ active_deadline.check()
+
# Headless environments: use authcode redirect flow (no local server)
use_authcode = _is_headless_environment() or os.environ.get(
"DROIDRUN_OAUTH_MANUAL", ""
@@ -512,11 +556,14 @@ def login(
open_browser=open_browser,
timeout_seconds=timeout_seconds,
prompt_consent=prompt_consent,
+ deadline=active_deadline,
+ persist_credentials=persist_credentials,
)
# Desktop: browser callback server
result: Dict[str, Optional[str]] = {"code": None, "state": None, "error": None}
done = threading.Event()
+ callback_lock = threading.Lock()
expected_state = secrets.token_hex(32)
code_verifier, code_challenge = _pkce_pair()
@@ -528,12 +575,17 @@ def do_GET(self) -> None: # noqa: N802
self.end_headers()
return
- params = parse_qs(parsed.query)
- result["code"] = params.get("code", [None])[0]
- result["state"] = params.get("state", [None])[0]
- result["error"] = params.get("error", [None])[0]
-
- ok = result["code"] is not None and result["error"] is None
+ with callback_lock:
+ if done.is_set():
+ self.send_response(409)
+ self.end_headers()
+ return
+ params = parse_qs(parsed.query)
+ result["code"] = params.get("code", [None])[0]
+ result["state"] = params.get("state", [None])[0]
+ result["error"] = params.get("error", [None])[0]
+ ok = result["code"] is not None and result["error"] is None
+ done.set()
self.send_response(200 if ok else 400)
self.send_header("Content-Type", "text/html; charset=utf-8")
self.end_headers()
@@ -545,7 +597,6 @@ def do_GET(self) -> None: # noqa: N802
self.wfile.write(
b"Login failed. Return to your terminal.
"
)
- done.set()
def log_message(self, format: str, *args: Any) -> None: # noqa: A003
return
@@ -561,8 +612,11 @@ def log_message(self, format: str, *args: Any) -> None: # noqa: A003
open_browser=open_browser,
timeout_seconds=timeout_seconds,
prompt_consent=prompt_consent,
+ deadline=active_deadline,
+ persist_credentials=persist_credentials,
)
+ active_deadline.check()
actual_port = httpd.server_address[1]
redirect_uri = f"http://127.0.0.1:{actual_port}{callback_path}"
auth_url = self._build_auth_url(
@@ -573,17 +627,19 @@ def log_message(self, format: str, *args: Any) -> None: # noqa: A003
)
server_thread = threading.Thread(target=httpd.serve_forever, daemon=True)
+ httpd.daemon_threads = True
server_thread.start()
try:
print(f"Open this URL to login:\n{auth_url}\n")
if open_browser:
- webbrowser.open(auth_url)
+ open_browser_async(auth_url, webbrowser.open)
- if not done.wait(timeout=timeout_seconds):
+ if not done.wait(timeout=active_deadline.remaining()):
raise TimeoutError(
"OAuth login timed out before callback was received."
)
+ active_deadline.check()
if result["error"]:
raise RuntimeError(f"OAuth callback returned error: {result['error']}")
@@ -594,12 +650,17 @@ def log_message(self, format: str, *args: Any) -> None: # noqa: A003
"OAuth callback did not include an authorization code."
)
- return self._exchange_authorization_code(
- result["code"], redirect_uri, code_verifier=code_verifier
+ access_token = self._exchange_authorization_code(
+ result["code"],
+ redirect_uri,
+ code_verifier=code_verifier,
+ deadline=active_deadline,
+ persist_credentials=persist_credentials,
)
finally:
httpd.shutdown()
httpd.server_close()
+ return access_token
def login_headless(
self,
@@ -608,8 +669,13 @@ def login_headless(
timeout_seconds: float = 300.0,
input_fn: Any = input,
prompt_consent: bool = True,
+ deadline: OAuthLoginDeadline | None = None,
+ persist_credentials: bool = True,
) -> str:
"""Headless OAuth flow for SSH/WSL environments."""
+ active_deadline = deadline or OAuthLoginDeadline(timeout_seconds)
+ active_deadline.check()
+
code_verifier, code_challenge = _pkce_pair()
expected_state = secrets.token_hex(32)
redirect_uri = "https://codeassist.google.com/authcode"
@@ -627,9 +693,8 @@ def login_headless(
f"\n2. Complete sign-in, then paste the authorization code shown on the page.\n"
)
if open_browser:
- webbrowser.open(auth_url)
+ open_browser_async(auth_url, webbrowser.open)
- deadline = time.time() + timeout_seconds
input_queue: queue.Queue[Optional[str]] = queue.Queue()
stop = threading.Event()
need_more = threading.Event()
@@ -650,16 +715,13 @@ def _reader() -> None:
try:
for attempt in range(_MAX_CODE_ATTEMPTS):
- remaining = deadline - time.time()
- if remaining <= 0:
- raise TimeoutError("OAuth login timed out.")
-
need_more.set()
try:
- raw = input_queue.get(timeout=remaining)
+ raw = input_queue.get(timeout=active_deadline.remaining())
except queue.Empty:
raise TimeoutError("OAuth login timed out.") from None
+ active_deadline.check()
if raw is None:
raise RuntimeError("Login failed — stdin closed.")
@@ -677,7 +739,11 @@ def _reader() -> None:
raise RuntimeError("Login failed.") from None
if code:
return self._exchange_authorization_code(
- code, redirect_uri, code_verifier=code_verifier
+ code,
+ redirect_uri,
+ code_verifier=code_verifier,
+ deadline=active_deadline,
+ persist_credentials=persist_credentials,
)
if attempt == 0:
print("Invalid code. Try again.")
diff --git a/mobilerun/agent/utils/oauth/grok_oauth_llm.py b/mobilerun/agent/utils/oauth/grok_oauth_llm.py
new file mode 100644
index 00000000..e143bdf3
--- /dev/null
+++ b/mobilerun/agent/utils/oauth/grok_oauth_llm.py
@@ -0,0 +1,1121 @@
+"""xAI subscription OAuth transport for the Responses API.
+
+Credentials are owned by Mobilerun and stored in its shared auth profile. The
+OAuth client and inference proxy values are deliberately pinned: accepting
+caller-controlled endpoints would let an attacker exfiltrate refresh tokens.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import base64
+import hashlib
+import os
+import re
+import secrets
+import sys
+import threading
+import time
+import webbrowser
+from dataclasses import dataclass
+from http.server import BaseHTTPRequestHandler
+from http.server import ThreadingHTTPServer as HTTPServer
+from pathlib import Path
+from typing import Any, Callable, Iterator
+from urllib.parse import parse_qs, urlencode, urlparse
+
+import httpx
+import jwt
+from llama_index.core.base.llms.types import LLMMetadata
+from llama_index.llms.openai.responses import OpenAIResponses
+from llama_index.llms.openai.utils import to_openai_message_dicts
+
+from mobilerun.agent.providers.grok import (
+ GROK_MODELS,
+ normalize_grok_model_id,
+ sanitize_grok_responses_kwargs,
+)
+from mobilerun.agent.utils.oauth.login_timeout import (
+ OAuthLoginDeadline,
+ open_browser_async,
+)
+from mobilerun.config_manager.auth_profile_store import AuthProfileStore
+from mobilerun.config_manager.credential_paths import GROK_OAUTH_CREDENTIAL_PATH
+
+DEFAULT_GROK_MODEL = "grok-4.5"
+DEFAULT_GROK_CONTEXT_WINDOW = 500_000
+DEFAULT_GROK_OAUTH_ISSUER = "https://auth.x.ai"
+DEFAULT_GROK_OAUTH_AUTHORIZE_URL = f"{DEFAULT_GROK_OAUTH_ISSUER}/oauth2/authorize"
+DEFAULT_GROK_OAUTH_DEVICE_URL = f"{DEFAULT_GROK_OAUTH_ISSUER}/oauth2/device/code"
+DEFAULT_GROK_OAUTH_TOKEN_URL = f"{DEFAULT_GROK_OAUTH_ISSUER}/oauth2/token"
+DEFAULT_GROK_OAUTH_JWKS_URL = f"{DEFAULT_GROK_OAUTH_ISSUER}/.well-known/jwks.json"
+DEFAULT_GROK_OAUTH_CLIENT_ID = "b1a00492-073a-47ea-816f-4c329264a828"
+DEFAULT_GROK_OAUTH_PROXY = "https://cli-chat-proxy.grok.com/v1"
+# The first-party Grok Build 1.0.0 proxy clients send this exact protocol
+# header/value. Keep it pinned rather than deriving it from any installed CLI.
+GROK_CLI_COMPAT_VERSION_HEADER = "x-grok-client-version"
+GROK_CLI_COMPAT_VERSION = "1.0.0"
+DEFAULT_GROK_OAUTH_CREDENTIAL_PATH = GROK_OAUTH_CREDENTIAL_PATH
+DEFAULT_GROK_OAUTH_SLOT = "grokOauth"
+DEFAULT_GROK_OAUTH_CALLBACK_HOST = "127.0.0.1"
+DEFAULT_GROK_OAUTH_CALLBACK_PORT = 0
+DEFAULT_GROK_OAUTH_CALLBACK_PATH = "/callback"
+DEFAULT_GROK_OAUTH_SCOPES = (
+ "openid",
+ "profile",
+ "email",
+ "offline_access",
+ "grok-cli:access",
+ "api:access",
+ "conversations:read",
+ "conversations:write",
+ "workspaces:read",
+ "workspaces:write",
+)
+DEFAULT_REFRESH_SKEW_SECONDS = 300
+DEFAULT_TOKEN_RETRY_BACKOFF_SECONDS = (0.25, 0.5)
+_DEVICE_GRANT = "urn:ietf:params:oauth:grant-type:device_code"
+_DEVICE_SURFACE_HEADER = "x-grok-client-surface"
+_DEVICE_SURFACE = "grok-build"
+_OAUTH_ERROR_CODE = re.compile(r"^[A-Za-z0-9_.-]{1,64}$")
+_TOKEN_ERROR_CODES = frozenset(
+ {
+ "access_denied",
+ "authorization_pending",
+ "expired_token",
+ "invalid_client",
+ "invalid_grant",
+ "invalid_request",
+ "invalid_scope",
+ "slow_down",
+ "unauthorized_client",
+ "unsupported_grant_type",
+ }
+)
+
+
+def _b64_no_pad(raw: bytes) -> str:
+ return base64.urlsafe_b64encode(raw).decode("ascii").rstrip("=")
+
+
+def _pkce_pair() -> tuple[str, str]:
+ verifier = _b64_no_pad(secrets.token_bytes(64))
+ challenge = _b64_no_pad(hashlib.sha256(verifier.encode("ascii")).digest())
+ return verifier, challenge
+
+
+def _is_headless_environment() -> bool:
+ if os.environ.get("SSH_CONNECTION") or os.environ.get("SSH_TTY"):
+ return True
+ if os.environ.get("WSL_DISTRO_NAME"):
+ return True
+ return sys.platform.startswith("linux") and not (
+ os.environ.get("DISPLAY") or os.environ.get("WAYLAND_DISPLAY")
+ )
+
+
+def _parse_callback_query(query: str) -> dict[str, str | None]:
+ """Parse a callback exactly once, rejecting ambiguous duplicate values."""
+ # Count blank values too: ``code=&code=value`` is still an ambiguous,
+ # duplicated security-sensitive parameter and must not be normalized into
+ # a single accepted value.
+ params = parse_qs(query, keep_blank_values=True)
+ code_values = params.get("code", [])
+ state_values = params.get("state", [])
+ error_values = params.get("error", [])
+ valid_cardinality = (
+ len(code_values) <= 1 and len(state_values) <= 1 and len(error_values) <= 1
+ )
+ return {
+ "code": code_values[0] if len(code_values) == 1 else None,
+ "state": state_values[0] if len(state_values) == 1 else None,
+ "error": (
+ error_values[0]
+ if len(error_values) == 1
+ else ("invalid_callback" if not valid_cardinality else None)
+ ),
+ }
+
+
+class GrokOAuthError(RuntimeError):
+ """A safe OAuth failure that never includes a token response body."""
+
+
+class GrokOAuthReloginRequired(GrokOAuthError):
+ """The refresh grant is permanently invalid and login must be repeated."""
+
+
+def _safe_error_code(value: object) -> str | None:
+ return (
+ value if isinstance(value, str) and _OAUTH_ERROR_CODE.fullmatch(value) else None
+ )
+
+
+def _safe_token_error_code(value: object) -> str | None:
+ return value if isinstance(value, str) and value in _TOKEN_ERROR_CODES else None
+
+
+def _safe_json_object(response: httpx.Response, *, context: str) -> dict[str, Any]:
+ try:
+ payload = response.json()
+ except Exception as exc:
+ raise GrokOAuthError(f"{context} did not return valid JSON.") from exc
+ if not isinstance(payload, dict):
+ raise GrokOAuthError(f"{context} did not return a JSON object.")
+ return payload
+
+
+@dataclass(frozen=True)
+class GrokOAuthCredentials:
+ access_token: str
+ refresh_token: str | None
+ expires_at_ms: int | None
+ token_type: str = "Bearer"
+ scopes: tuple[str, ...] = DEFAULT_GROK_OAUTH_SCOPES
+ issuer: str = DEFAULT_GROK_OAUTH_ISSUER
+ client_id: str = DEFAULT_GROK_OAUTH_CLIENT_ID
+
+ @classmethod
+ def from_payload(cls, payload: dict[str, Any]) -> "GrokOAuthCredentials":
+ if payload.get("type") != "oauth" or payload.get("provider") != "xai-grok":
+ raise ValueError("XAI OAuth profile has an unexpected credential type.")
+ access_token = payload.get("accessToken")
+ if not isinstance(access_token, str) or not access_token:
+ raise ValueError("XAI OAuth profile is missing accessToken.")
+ issuer = payload.get("issuer")
+ client_id = payload.get("clientId")
+ if issuer != DEFAULT_GROK_OAUTH_ISSUER:
+ raise ValueError("XAI OAuth profile has an unexpected issuer.")
+ if client_id != DEFAULT_GROK_OAUTH_CLIENT_ID:
+ raise ValueError("XAI OAuth profile has an unexpected clientId.")
+
+ refresh = payload.get("refreshToken")
+ raw_expiry = payload.get("expiresAt")
+ try:
+ expires_at_ms = int(raw_expiry) if raw_expiry is not None else None
+ except (TypeError, ValueError):
+ expires_at_ms = None
+ raw_scopes = payload.get("scopes")
+ scopes = (
+ tuple(str(value) for value in raw_scopes if isinstance(value, str))
+ if isinstance(raw_scopes, list)
+ else DEFAULT_GROK_OAUTH_SCOPES
+ )
+ token_type = str(payload.get("tokenType") or "Bearer")
+ if token_type.lower() != "bearer":
+ raise ValueError("XAI OAuth profile has an unsupported tokenType.")
+ return cls(
+ access_token=access_token,
+ refresh_token=refresh if isinstance(refresh, str) and refresh else None,
+ expires_at_ms=expires_at_ms,
+ token_type="Bearer",
+ scopes=scopes,
+ issuer=issuer,
+ client_id=client_id,
+ )
+
+ def to_payload(self) -> dict[str, Any]:
+ return {
+ "type": "oauth",
+ "provider": "xai-grok",
+ "accessToken": self.access_token,
+ "refreshToken": self.refresh_token,
+ "expiresAt": self.expires_at_ms,
+ "tokenType": self.token_type,
+ "scopes": list(self.scopes),
+ "issuer": self.issuer,
+ "clientId": self.client_id,
+ }
+
+ def is_valid(self, *, skew_ms: int = DEFAULT_REFRESH_SKEW_SECONDS * 1000) -> bool:
+ return self.expires_at_ms is None or (
+ int(time.time() * 1000) + skew_ms < self.expires_at_ms
+ )
+
+
+class GrokOAuthCredentialStore:
+ def __init__(self, path: str | Path = DEFAULT_GROK_OAUTH_CREDENTIAL_PATH) -> None:
+ self.path = Path(path).expanduser()
+ self.profile_store = AuthProfileStore(self.path)
+
+ def load(self) -> GrokOAuthCredentials | None:
+ payload = self.profile_store.read_slot(DEFAULT_GROK_OAUTH_SLOT)
+ if payload is None:
+ return None
+ return GrokOAuthCredentials.from_payload(payload)
+
+ def save(
+ self,
+ credentials: GrokOAuthCredentials,
+ *,
+ deadline: OAuthLoginDeadline | None = None,
+ ) -> None:
+ if deadline is not None:
+ deadline.check()
+ self.profile_store.update_slot(
+ DEFAULT_GROK_OAUTH_SLOT,
+ credentials.to_payload(),
+ lock_timeout=deadline.remaining() if deadline is not None else None,
+ before_commit=deadline.check if deadline is not None else None,
+ )
+
+
+class GrokIDTokenValidator:
+ """Validate xAI ID tokens using its pinned ES256 JWKS endpoint."""
+
+ def __init__(self, jwks_client: Any | None = None) -> None:
+ self._jwks_client = jwks_client or jwt.PyJWKClient(DEFAULT_GROK_OAUTH_JWKS_URL)
+ self._jwks_lock = threading.Lock()
+
+ def validate(
+ self,
+ token: str,
+ *,
+ nonce: str | None,
+ deadline: OAuthLoginDeadline | None = None,
+ ) -> dict[str, Any]:
+ if deadline is None:
+ self._jwks_lock.acquire()
+ elif not self._jwks_lock.acquire(timeout=deadline.remaining()):
+ raise TimeoutError("XAI OAuth login timed out.")
+ had_timeout = hasattr(self._jwks_client, "timeout")
+ original_timeout = getattr(self._jwks_client, "timeout", None)
+ timeout_changed = False
+ instance_attributes = getattr(self._jwks_client, "__dict__", {})
+ had_fetch_override = "fetch_data" in instance_attributes
+ original_fetch_override = instance_attributes.get("fetch_data")
+ fetch_changed = False
+
+ def _remaining_request_timeout() -> float:
+ if deadline is None:
+ raise AssertionError("JWKS deadline wrapper requires a deadline")
+ try:
+ return deadline.remaining(cap=float(original_timeout))
+ except (TypeError, ValueError):
+ return deadline.remaining()
+
+ try:
+ if deadline is not None:
+ deadline.check()
+ if type(self._jwks_client) is jwt.PyJWKClient:
+ original_fetch = self._jwks_client.fetch_data
+
+ def _fetch_with_remaining_deadline() -> Any:
+ self._jwks_client.timeout = _remaining_request_timeout()
+ return original_fetch()
+
+ self._jwks_client.fetch_data = _fetch_with_remaining_deadline
+ fetch_changed = True
+ timeout_changed = True
+ else:
+ try:
+ self._jwks_client.timeout = _remaining_request_timeout()
+ timeout_changed = True
+ except (AttributeError, TypeError, ValueError):
+ # Arbitrary injected validators may not expose a writable
+ # network-timeout setting; deadline checks still bracket
+ # the call without breaking their existing contract.
+ pass
+ signing_key = self._jwks_client.get_signing_key_from_jwt(token).key
+ if deadline is not None:
+ deadline.check()
+ finally:
+ if fetch_changed:
+ if had_fetch_override:
+ self._jwks_client.fetch_data = original_fetch_override
+ else:
+ del self._jwks_client.fetch_data
+ if timeout_changed:
+ try:
+ if had_timeout:
+ self._jwks_client.timeout = original_timeout
+ else:
+ del self._jwks_client.timeout
+ except (AttributeError, TypeError):
+ pass
+ self._jwks_lock.release()
+ claims = jwt.decode(
+ token,
+ signing_key,
+ algorithms=["ES256"],
+ issuer=DEFAULT_GROK_OAUTH_ISSUER,
+ audience=DEFAULT_GROK_OAUTH_CLIENT_ID,
+ options={"require": ["exp", "iat", "iss", "aud", "sub"]},
+ )
+ if nonce is not None and claims.get("nonce") != nonce:
+ raise jwt.InvalidTokenError("ID token nonce mismatch.")
+ if deadline is not None:
+ deadline.check()
+ return claims
+
+
+class GrokOAuthSessionManager:
+ """Load, exchange, and cross-process-refresh Mobilerun's xAI session."""
+
+ def __init__(
+ self,
+ *,
+ credential_store: GrokOAuthCredentialStore | None = None,
+ http_client: httpx.Client | None = None,
+ id_token_validator: GrokIDTokenValidator | None = None,
+ request_timeout: float = 20.0,
+ refresh_skew_seconds: int = DEFAULT_REFRESH_SKEW_SECONDS,
+ retry_backoff_seconds: tuple[float, ...] = DEFAULT_TOKEN_RETRY_BACKOFF_SECONDS,
+ sleep: Callable[[float], None] | None = None,
+ ) -> None:
+ self.credential_store = credential_store or GrokOAuthCredentialStore()
+ self.http_client = http_client or httpx.Client()
+ self.id_token_validator = id_token_validator or GrokIDTokenValidator()
+ self.request_timeout = request_timeout
+ self.refresh_skew_ms = max(0, int(refresh_skew_seconds)) * 1000
+ self.retry_backoff_seconds = tuple(
+ max(0.0, float(delay)) for delay in retry_backoff_seconds[:2]
+ )
+ self.sleep = sleep or time.sleep
+ self._thread_lock = threading.RLock()
+ self._credentials: GrokOAuthCredentials | None = None
+
+ @staticmethod
+ def _expiry_ms(payload: dict[str, Any]) -> int | None:
+ try:
+ return int(time.time() * 1000) + int(payload["expires_in"]) * 1000
+ except (KeyError, TypeError, ValueError):
+ pass
+
+ # Expiry is scheduling metadata, not an authorization decision. If
+ # the provider omits expires_in, the JWT exp claim prevents treating a
+ # short-lived access token as permanent; signature validation remains
+ # mandatory for ID-token identity claims.
+ for key in ("id_token", "access_token"):
+ token = payload.get(key)
+ if not isinstance(token, str):
+ continue
+ try:
+ claims = jwt.decode(
+ token,
+ options={"verify_signature": False, "verify_exp": False},
+ algorithms=["ES256"],
+ )
+ return int(claims["exp"]) * 1000
+ except (jwt.PyJWTError, KeyError, TypeError, ValueError):
+ continue
+ return None
+
+ @staticmethod
+ def _scopes(payload: dict[str, Any], fallback: tuple[str, ...]) -> tuple[str, ...]:
+ raw_scope = payload.get("scope")
+ if isinstance(raw_scope, str) and raw_scope.strip():
+ return tuple(raw_scope.split())
+ return fallback
+
+ def _credentials_from_token_response(
+ self,
+ payload: dict[str, Any],
+ *,
+ prior_refresh_token: str | None = None,
+ prior_scopes: tuple[str, ...] = DEFAULT_GROK_OAUTH_SCOPES,
+ nonce: str | None = None,
+ deadline: OAuthLoginDeadline | None = None,
+ ) -> GrokOAuthCredentials:
+ if deadline is not None:
+ deadline.check()
+ access_token = payload.get("access_token")
+ if not isinstance(access_token, str) or not access_token:
+ raise GrokOAuthError("xAI token response did not contain an access token.")
+ id_token = payload.get("id_token")
+ if nonce is not None and not (isinstance(id_token, str) and id_token):
+ raise GrokOAuthError(
+ "xAI authorization response did not contain an ID token."
+ )
+ if isinstance(id_token, str) and id_token:
+ if (
+ deadline is None
+ or type(self.id_token_validator) is not GrokIDTokenValidator
+ ):
+ if deadline is not None:
+ deadline.check()
+ self.id_token_validator.validate(id_token, nonce=nonce)
+ if deadline is not None:
+ deadline.check()
+ else:
+ self.id_token_validator.validate(
+ id_token,
+ nonce=nonce,
+ deadline=deadline,
+ )
+
+ refresh_token = payload.get("refresh_token")
+ if not isinstance(refresh_token, str) or not refresh_token:
+ refresh_token = prior_refresh_token
+ token_type = payload.get("token_type")
+ if isinstance(token_type, str) and token_type.lower() != "bearer":
+ raise GrokOAuthError("xAI token response used an unsupported token type.")
+ credentials = GrokOAuthCredentials(
+ access_token=access_token,
+ refresh_token=refresh_token,
+ expires_at_ms=self._expiry_ms(payload),
+ token_type="Bearer",
+ scopes=self._scopes(payload, prior_scopes),
+ )
+ if deadline is not None:
+ deadline.check()
+ return credentials
+
+ def _post_form(
+ self,
+ url: str,
+ *,
+ data: dict[str, str],
+ headers: dict[str, str],
+ context: str,
+ retry_transient: bool,
+ deadline: OAuthLoginDeadline | None = None,
+ ) -> httpx.Response:
+ """POST a form, retrying only grants that are safe to replay."""
+ backoffs = self.retry_backoff_seconds if retry_transient else ()
+ for attempt in range(len(backoffs) + 1):
+ request_timeout = (
+ deadline.remaining(cap=self.request_timeout)
+ if deadline is not None
+ else self.request_timeout
+ )
+ try:
+ response = self.http_client.post(
+ url,
+ headers=headers,
+ data=data,
+ timeout=request_timeout,
+ )
+ except (httpx.ConnectError, httpx.TimeoutException) as exc:
+ if deadline is not None:
+ deadline.check()
+ if attempt < len(backoffs):
+ if deadline is not None:
+ deadline.sleep(backoffs[attempt])
+ else:
+ self.sleep(backoffs[attempt])
+ continue
+ raise GrokOAuthError(
+ f"{context} failed due to a transient network error."
+ ) from exc
+ if deadline is not None:
+ deadline.check()
+ if 500 <= response.status_code < 600 and attempt < len(backoffs):
+ if deadline is not None:
+ deadline.sleep(backoffs[attempt])
+ else:
+ self.sleep(backoffs[attempt])
+ continue
+ return response
+ raise AssertionError("unreachable token retry state")
+
+ def _post_token(
+ self,
+ data: dict[str, str],
+ *,
+ retry_transient: bool = False,
+ refresh_request: bool = False,
+ deadline: OAuthLoginDeadline | None = None,
+ ) -> dict[str, Any]:
+ response = self._post_form(
+ DEFAULT_GROK_OAUTH_TOKEN_URL,
+ headers={
+ "Accept": "application/json",
+ "Content-Type": "application/x-www-form-urlencoded",
+ },
+ data=data,
+ context="xAI token request",
+ retry_transient=retry_transient,
+ deadline=deadline,
+ )
+ if response.status_code >= 400:
+ try:
+ error = _safe_token_error_code(response.json().get("error"))
+ except Exception:
+ error = None
+ if refresh_request and error in {"invalid_grant", "invalid_client"}:
+ raise GrokOAuthReloginRequired(
+ "XAI OAuth refresh was rejected; re-login is required."
+ )
+ raise GrokOAuthError(
+ f"xAI token request failed ({error or response.status_code})."
+ )
+ payload = _safe_json_object(response, context="xAI token response")
+ if deadline is not None:
+ deadline.check()
+ return payload
+
+ def set_initial_credentials(
+ self,
+ credentials: GrokOAuthCredentials,
+ *,
+ deadline: OAuthLoginDeadline | None = None,
+ ) -> None:
+ if deadline is None:
+ with self._thread_lock:
+ self.credential_store.save(credentials)
+ self._credentials = credentials
+ return
+
+ deadline.check()
+ if not self._thread_lock.acquire(timeout=deadline.remaining()):
+ raise TimeoutError("XAI OAuth login timed out.")
+ try:
+ self.credential_store.save(credentials, deadline=deadline)
+ self._credentials = credentials
+ finally:
+ self._thread_lock.release()
+
+ def exchange_authorization_code(
+ self,
+ *,
+ code: str,
+ redirect_uri: str,
+ code_verifier: str,
+ nonce: str,
+ deadline: OAuthLoginDeadline | None = None,
+ ) -> GrokOAuthCredentials:
+ payload = self._post_token(
+ {
+ "grant_type": "authorization_code",
+ "client_id": DEFAULT_GROK_OAUTH_CLIENT_ID,
+ "code": code,
+ "redirect_uri": redirect_uri,
+ "code_verifier": code_verifier,
+ },
+ deadline=deadline,
+ )
+ credentials = self._credentials_from_token_response(
+ payload,
+ nonce=nonce,
+ deadline=deadline,
+ )
+ self.set_initial_credentials(credentials, deadline=deadline)
+ return credentials
+
+ def _refresh(self, credentials: GrokOAuthCredentials) -> GrokOAuthCredentials:
+ if not credentials.refresh_token:
+ raise ValueError(
+ "No XAI OAuth refresh token is available. Run `mobilerun configure xai`."
+ )
+ payload = self._post_token(
+ {
+ "grant_type": "refresh_token",
+ "client_id": DEFAULT_GROK_OAUTH_CLIENT_ID,
+ "refresh_token": credentials.refresh_token,
+ },
+ retry_transient=True,
+ refresh_request=True,
+ )
+ return self._credentials_from_token_response(
+ payload,
+ prior_refresh_token=credentials.refresh_token,
+ prior_scopes=credentials.scopes,
+ )
+
+ def get_valid_credentials(
+ self,
+ *,
+ force_refresh: bool = False,
+ rejected_access_token: str | None = None,
+ ) -> GrokOAuthCredentials:
+ with self._thread_lock:
+ # Keep the file lock across refresh so separate Mobilerun processes
+ # cannot rotate the same refresh token concurrently.
+ with self.credential_store.profile_store.transaction() as transaction:
+ payload = transaction.get_slot(DEFAULT_GROK_OAUTH_SLOT)
+ credentials = (
+ GrokOAuthCredentials.from_payload(payload)
+ if payload is not None
+ else self._credentials
+ )
+ if credentials is None:
+ raise ValueError(
+ "No XAI OAuth credentials found. Run `mobilerun configure xai`."
+ )
+
+ another_writer_refreshed = (
+ rejected_access_token is not None
+ and credentials.access_token != rejected_access_token
+ and credentials.is_valid(skew_ms=self.refresh_skew_ms)
+ )
+ if another_writer_refreshed or (
+ not force_refresh
+ and credentials.is_valid(skew_ms=self.refresh_skew_ms)
+ ):
+ self._credentials = credentials
+ return credentials
+
+ try:
+ refreshed = self._refresh(credentials)
+ except GrokOAuthReloginRequired:
+ # Do not keep using a bearer token whose refresh grant is
+ # permanently invalid. The locked file remains untouched,
+ # including every sibling provider slot.
+ self._credentials = None
+ raise
+ transaction.set_slot(DEFAULT_GROK_OAUTH_SLOT, refreshed.to_payload())
+ self._credentials = refreshed
+ return refreshed
+
+
+class GrokOAuthAuth(httpx.Auth):
+ """Inject a fresh bearer token and retry one rejected request once."""
+
+ requires_request_body = True
+
+ def __init__(self, manager: GrokOAuthSessionManager, *, model: str) -> None:
+ self.manager = manager
+ self.model = model
+
+ def _authorize(
+ self, request: httpx.Request, credentials: GrokOAuthCredentials
+ ) -> None:
+ request.headers["Authorization"] = (
+ f"{credentials.token_type} {credentials.access_token}"
+ )
+ request.headers["X-XAI-Token-Auth"] = "xai-grok-cli"
+ request.headers["x-grok-model-override"] = self.model
+ request.headers[GROK_CLI_COMPAT_VERSION_HEADER] = GROK_CLI_COMPAT_VERSION
+
+ def sync_auth_flow(self, request: httpx.Request) -> Iterator[httpx.Request]:
+ credentials = self.manager.get_valid_credentials()
+ self._authorize(request, credentials)
+ response = yield request
+ if response.status_code == 401:
+ response.read()
+ credentials = self.manager.get_valid_credentials(
+ force_refresh=True,
+ rejected_access_token=credentials.access_token,
+ )
+ self._authorize(request, credentials)
+ yield request
+
+ async def async_auth_flow(self, request: httpx.Request): # type: ignore[no-untyped-def]
+ credentials = await asyncio.to_thread(self.manager.get_valid_credentials)
+ self._authorize(request, credentials)
+ response = yield request
+ if response.status_code == 401:
+ await response.aread()
+ credentials = await asyncio.to_thread(
+ self.manager.get_valid_credentials,
+ force_refresh=True,
+ rejected_access_token=credentials.access_token,
+ )
+ self._authorize(request, credentials)
+ yield request
+
+
+class GrokOAuth(OpenAIResponses):
+ """LlamaIndex Responses adapter backed only by Mobilerun's xAI OAuth slot."""
+
+ @classmethod
+ def class_name(cls) -> str:
+ return "GrokOAuth"
+
+ def __init__(
+ self,
+ model: str = DEFAULT_GROK_MODEL,
+ oauth_credential_path: str | None = None,
+ credential_path: str | None = None,
+ oauth_access_token: str | None = None,
+ oauth_refresh_token: str | None = None,
+ oauth_expires_at_ms: int | None = None,
+ oauth_refresh_skew_seconds: int = DEFAULT_REFRESH_SKEW_SECONDS,
+ oauth_session_manager: GrokOAuthSessionManager | None = None,
+ http_client: httpx.Client | None = None,
+ async_http_client: httpx.AsyncClient | None = None,
+ **kwargs: Any,
+ ) -> None:
+ model = normalize_grok_model_id(model)
+ if model not in GROK_MODELS:
+ raise ValueError(
+ f"Model {model!r} is not supported with XAI OAuth. "
+ f"Use {', '.join(GROK_MODELS)}."
+ )
+ path = (
+ oauth_credential_path
+ or credential_path
+ or str(DEFAULT_GROK_OAUTH_CREDENTIAL_PATH)
+ )
+ manager = oauth_session_manager or GrokOAuthSessionManager(
+ credential_store=GrokOAuthCredentialStore(path),
+ request_timeout=float(kwargs.get("timeout", 60.0)),
+ refresh_skew_seconds=oauth_refresh_skew_seconds,
+ )
+ if oauth_access_token or oauth_refresh_token:
+ manager.set_initial_credentials(
+ GrokOAuthCredentials(
+ access_token=oauth_access_token or "oauth",
+ refresh_token=oauth_refresh_token,
+ expires_at_ms=(oauth_expires_at_ms if oauth_access_token else 0),
+ )
+ )
+
+ auth = GrokOAuthAuth(manager, model=model)
+ if http_client is None:
+ http_client = httpx.Client(auth=auth)
+ else:
+ http_client.auth = auth
+ if async_http_client is None:
+ async_http_client = httpx.AsyncClient(auth=auth)
+ else:
+ async_http_client.auth = auth
+
+ supplied_headers = dict(kwargs.pop("default_headers", None) or {})
+ supplied_headers.update(
+ {
+ "X-XAI-Token-Auth": "xai-grok-cli",
+ "x-grok-model-override": model,
+ GROK_CLI_COMPAT_VERSION_HEADER: GROK_CLI_COMPAT_VERSION,
+ }
+ )
+ kwargs.pop("api_key", None)
+ kwargs.pop("api_base", None)
+ kwargs.pop("base_url", None)
+ kwargs.pop("store", None)
+ kwargs.pop("track_previous_responses", None)
+ kwargs.pop("reasoning_options", None)
+ kwargs.pop("context_window", None)
+ kwargs.pop("openai_client", None)
+ kwargs.pop("async_openai_client", None)
+ super().__init__(
+ model=model,
+ api_key="oauth",
+ api_base=DEFAULT_GROK_OAUTH_PROXY,
+ context_window=DEFAULT_GROK_CONTEXT_WINDOW,
+ store=False,
+ track_previous_responses=False,
+ reasoning_options=None,
+ default_headers=supplied_headers,
+ http_client=http_client,
+ async_http_client=async_http_client,
+ **kwargs,
+ )
+ self._oauth_manager = manager
+
+ @property
+ def metadata(self) -> LLMMetadata:
+ return LLMMetadata(
+ context_window=DEFAULT_GROK_CONTEXT_WINDOW,
+ num_output=self.max_output_tokens or -1,
+ is_chat_model=True,
+ is_function_calling_model=True,
+ model_name=self.model,
+ )
+
+ @property
+ def _tokenizer(self): # type: ignore[no-untyped-def]
+ # tiktoken does not have an encoding registered for xAI model ids.
+ return None
+
+ def _get_model_kwargs(self, **kwargs: Any) -> dict[str, Any]:
+ model_kwargs = super()._get_model_kwargs(**kwargs)
+ return self._sanitize_request_kwargs(model_kwargs)
+
+ def _sanitize_request_kwargs(
+ self,
+ model_kwargs: dict[str, Any],
+ *,
+ omit_tool_choice: bool = False,
+ ) -> dict[str, Any]:
+ sanitized = dict(
+ sanitize_grok_responses_kwargs(
+ dict(model_kwargs),
+ omit_sampler_fields=True,
+ omit_tool_choice=omit_tool_choice,
+ )
+ )
+ # Both normal kwargs and ``extra_body`` are merged after constructor
+ # defaults. Pin the payload model as well as the proxy override header.
+ sanitized["model"] = self.model
+ return sanitized
+
+ def structured_predict(
+ self,
+ output_cls: Any,
+ prompt: Any,
+ llm_kwargs: dict[str, Any] | None = None,
+ **prompt_args: Any,
+ ) -> Any:
+ sanitized = self._sanitize_request_kwargs(
+ dict(llm_kwargs or {}), omit_tool_choice=True
+ )
+ sanitized.pop("model", None)
+ sanitized.pop("tool_choice", None)
+ messages = prompt.format_messages(**prompt_args)
+ message_dicts = to_openai_message_dicts(
+ messages, model=self.model, is_responses_api=True
+ )
+ response = self._client.responses.parse(
+ model=self._responses_model,
+ input=message_dicts,
+ text_format=output_cls,
+ **sanitized,
+ )
+ if response.output_parsed is not None:
+ return response.output_parsed
+ raise ValueError("Failed to produce a structured response from the model.")
+
+ async def astructured_predict(
+ self,
+ output_cls: Any,
+ prompt: Any,
+ llm_kwargs: dict[str, Any] | None = None,
+ **prompt_args: Any,
+ ) -> Any:
+ sanitized = self._sanitize_request_kwargs(
+ dict(llm_kwargs or {}), omit_tool_choice=True
+ )
+ sanitized.pop("model", None)
+ sanitized.pop("tool_choice", None)
+ messages = prompt.format_messages(**prompt_args)
+ message_dicts = to_openai_message_dicts(
+ messages, model=self.model, is_responses_api=True
+ )
+ response = await self._aclient.responses.parse(
+ model=self._responses_model,
+ input=message_dicts,
+ text_format=output_cls,
+ **sanitized,
+ )
+ if response.output_parsed is not None:
+ return response.output_parsed
+ raise ValueError("Failed to produce a structured response from the model.")
+
+ @staticmethod
+ def _build_auth_url(
+ *, redirect_uri: str, code_challenge: str, state: str, nonce: str
+ ) -> str:
+ return f"{DEFAULT_GROK_OAUTH_AUTHORIZE_URL}?{urlencode({
+ 'response_type': 'code',
+ 'client_id': DEFAULT_GROK_OAUTH_CLIENT_ID,
+ 'redirect_uri': redirect_uri,
+ 'scope': ' '.join(DEFAULT_GROK_OAUTH_SCOPES),
+ 'code_challenge': code_challenge,
+ 'code_challenge_method': 'S256',
+ 'state': state,
+ 'nonce': nonce,
+ 'referrer': 'grok-build',
+ })}"
+
+ def login(
+ self,
+ *,
+ open_browser: bool = True,
+ timeout_seconds: float = 300.0,
+ callback_host: str = DEFAULT_GROK_OAUTH_CALLBACK_HOST,
+ callback_port: int = DEFAULT_GROK_OAUTH_CALLBACK_PORT,
+ callback_path: str = DEFAULT_GROK_OAUTH_CALLBACK_PATH,
+ device_code: bool = False,
+ deadline: OAuthLoginDeadline | None = None,
+ ) -> GrokOAuthCredentials:
+ login_deadline = deadline or OAuthLoginDeadline(
+ timeout_seconds,
+ timeout_message="XAI OAuth login timed out.",
+ sleeper=self._oauth_manager.sleep,
+ )
+ if callback_host != DEFAULT_GROK_OAUTH_CALLBACK_HOST:
+ raise ValueError("XAI OAuth callback_host must be 127.0.0.1.")
+ if callback_port != 0:
+ raise ValueError("XAI OAuth callback_port must be OS-assigned (0).")
+ if callback_path != DEFAULT_GROK_OAUTH_CALLBACK_PATH:
+ raise ValueError("XAI OAuth callback_path must be /callback.")
+ if device_code or _is_headless_environment():
+ return self._login_device_code(
+ deadline=login_deadline,
+ open_browser=open_browser,
+ )
+
+ code_verifier, code_challenge = _pkce_pair()
+ state = _b64_no_pad(secrets.token_bytes(32))
+ nonce = _b64_no_pad(secrets.token_bytes(32))
+ result: dict[str, str | None] = {
+ "code": None,
+ "state": None,
+ "error": None,
+ }
+ done = threading.Event()
+ callback_lock = threading.Lock()
+
+ class _CallbackHandler(BaseHTTPRequestHandler):
+ def do_GET(self) -> None: # noqa: N802
+ parsed = urlparse(self.path)
+ if parsed.path != DEFAULT_GROK_OAUTH_CALLBACK_PATH:
+ self.send_response(404)
+ self.end_headers()
+ return
+ with callback_lock:
+ if done.is_set():
+ self.send_response(409)
+ self.end_headers()
+ return
+ result.update(_parse_callback_query(parsed.query))
+ ok = bool(result["code"] and not result["error"])
+ # Callback receipt completes the wait. Do not let a stalled
+ # browser socket consume the remaining login deadline.
+ done.set()
+ self.send_response(200 if ok else 400)
+ self.send_header("Content-Type", "text/html; charset=utf-8")
+ self.end_headers()
+ self.wfile.write(
+ b"Mobilerun XAI login complete. You may close this tab."
+ if ok
+ else b"Mobilerun XAI login failed. Return to the terminal."
+ )
+
+ def log_message(self, format: str, *args: Any) -> None: # noqa: A003
+ return
+
+ try:
+ server = HTTPServer((DEFAULT_GROK_OAUTH_CALLBACK_HOST, 0), _CallbackHandler)
+ except OSError:
+ return self._login_device_code(
+ deadline=login_deadline,
+ open_browser=open_browser,
+ )
+
+ redirect_uri = (
+ f"http://{DEFAULT_GROK_OAUTH_CALLBACK_HOST}:{server.server_address[1]}"
+ f"{DEFAULT_GROK_OAUTH_CALLBACK_PATH}"
+ )
+ authorization_url = self._build_auth_url(
+ redirect_uri=redirect_uri,
+ code_challenge=code_challenge,
+ state=state,
+ nonce=nonce,
+ )
+ thread = threading.Thread(target=server.serve_forever, daemon=True)
+ server.daemon_threads = True
+ thread.start()
+ try:
+ print(f"Open this URL to sign in to xAI:\n{authorization_url}\n")
+ if open_browser:
+ open_browser_async(authorization_url, webbrowser.open)
+ if not done.wait(timeout=login_deadline.remaining()):
+ raise TimeoutError("XAI OAuth login timed out waiting for callback.")
+ login_deadline.check()
+ if result["error"]:
+ raise GrokOAuthError(
+ f"xAI authorization failed ({_safe_error_code(result['error']) or 'oauth_error'})."
+ )
+ if not secrets.compare_digest(result["state"] or "", state):
+ raise GrokOAuthError("XAI OAuth callback state mismatch.")
+ if not result["code"]:
+ raise GrokOAuthError("XAI OAuth callback did not contain a code.")
+ return self._oauth_manager.exchange_authorization_code(
+ code=result["code"],
+ redirect_uri=redirect_uri,
+ code_verifier=code_verifier,
+ nonce=nonce,
+ deadline=login_deadline,
+ )
+ finally:
+ server.shutdown()
+ server.server_close()
+
+ def _login_device_code(
+ self,
+ *,
+ deadline: OAuthLoginDeadline,
+ open_browser: bool,
+ ) -> GrokOAuthCredentials:
+ manager = self._oauth_manager
+ response = manager._post_form(
+ DEFAULT_GROK_OAUTH_DEVICE_URL,
+ headers={
+ "Accept": "application/json",
+ "Content-Type": "application/x-www-form-urlencoded",
+ _DEVICE_SURFACE_HEADER: _DEVICE_SURFACE,
+ },
+ data={
+ "client_id": DEFAULT_GROK_OAUTH_CLIENT_ID,
+ "scope": " ".join(DEFAULT_GROK_OAUTH_SCOPES),
+ },
+ context="xAI device authorization request",
+ retry_transient=False,
+ deadline=deadline,
+ )
+ if response.status_code >= 400:
+ raise GrokOAuthError(
+ f"xAI device authorization failed ({response.status_code})."
+ )
+ payload = _safe_json_object(
+ response, context="xAI device authorization response"
+ )
+ device_code = payload.get("device_code")
+ user_code = payload.get("user_code")
+ verification_uri = payload.get("verification_uri_complete") or payload.get(
+ "verification_uri"
+ )
+ if not all(
+ isinstance(value, str) and value
+ for value in (device_code, user_code, verification_uri)
+ ):
+ raise GrokOAuthError("xAI device authorization response was incomplete.")
+ try:
+ expires_in = max(1, int(payload.get("expires_in", 1800)))
+ interval = max(1, int(payload.get("interval", 5)))
+ except (TypeError, ValueError):
+ expires_in, interval = 1800, 5
+ device_deadline = deadline.limited_to(expires_in)
+ print(
+ "Open this URL to sign in to xAI:\n"
+ f"{verification_uri}\n\nEnter code: {user_code}\n"
+ "Never share this device code.\n"
+ )
+ if open_browser:
+ open_browser_async(verification_uri, webbrowser.open)
+
+ while True:
+ device_deadline.check()
+ token_response = manager._post_form(
+ DEFAULT_GROK_OAUTH_TOKEN_URL,
+ headers={
+ "Accept": "application/json",
+ "Content-Type": "application/x-www-form-urlencoded",
+ },
+ data={
+ "grant_type": _DEVICE_GRANT,
+ "device_code": device_code,
+ "client_id": DEFAULT_GROK_OAUTH_CLIENT_ID,
+ },
+ context="xAI device token request",
+ retry_transient=True,
+ deadline=device_deadline,
+ )
+ if token_response.status_code < 400:
+ token_payload = _safe_json_object(
+ token_response, context="xAI device token response"
+ )
+ device_deadline.check()
+ credentials = manager._credentials_from_token_response(
+ token_payload,
+ deadline=device_deadline,
+ )
+ manager.set_initial_credentials(
+ credentials,
+ deadline=device_deadline,
+ )
+ return credentials
+ try:
+ error = _safe_token_error_code(token_response.json().get("error"))
+ except Exception:
+ error = None
+ if error == "authorization_pending":
+ pass
+ elif error == "slow_down":
+ interval += 5
+ elif error == "access_denied":
+ raise GrokOAuthError("xAI device authorization was denied.")
+ elif error == "expired_token":
+ raise TimeoutError("xAI device authorization expired.")
+ else:
+ raise GrokOAuthError(
+ f"xAI device token request failed ({error or token_response.status_code})."
+ )
+ device_deadline.sleep(interval)
+
+
+# Descriptive alias for callers that prefer the full class name.
+GrokOAuthLLM = GrokOAuth
diff --git a/mobilerun/agent/utils/oauth/login_timeout.py b/mobilerun/agent/utils/oauth/login_timeout.py
new file mode 100644
index 00000000..55b2b0d6
--- /dev/null
+++ b/mobilerun/agent/utils/oauth/login_timeout.py
@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+import math
+import threading
+import time
+from collections.abc import Callable
+
+
+class OAuthLoginDeadline:
+ """One monotonic timeout budget shared by every OAuth login stage."""
+
+ def __init__(
+ self,
+ timeout_seconds: float,
+ *,
+ timeout_message: str = "OAuth login timed out.",
+ clock: Callable[[], float] | None = None,
+ sleeper: Callable[[float], None] | None = None,
+ _expires_at: float | None = None,
+ ) -> None:
+ timeout = float(timeout_seconds)
+ if not math.isfinite(timeout) or timeout <= 0:
+ raise ValueError("OAuth login timeout must be a finite positive number.")
+ self._clock = clock or time.monotonic
+ self._sleeper = sleeper or time.sleep
+ self._timeout_message = timeout_message
+ self._expires_at = (
+ self._clock() + timeout if _expires_at is None else _expires_at
+ )
+
+ @property
+ def expires_at(self) -> float:
+ return self._expires_at
+
+ def remaining(self, *, cap: float | None = None) -> float:
+ remaining = self._expires_at - self._clock()
+ if remaining <= 0:
+ raise TimeoutError(self._timeout_message)
+ if cap is None:
+ return remaining
+ maximum = float(cap)
+ if not math.isfinite(maximum) or maximum <= 0:
+ raise ValueError("OAuth request timeout cap must be finite and positive.")
+ return min(remaining, maximum)
+
+ def check(self) -> None:
+ self.remaining()
+
+ def sleep(self, delay_seconds: float) -> None:
+ delay = max(0.0, float(delay_seconds))
+ if delay:
+ self._sleeper(min(delay, self.remaining()))
+ self.check()
+
+ def limited_to(self, timeout_seconds: float) -> OAuthLoginDeadline:
+ """Return a view capped by a provider-issued expiry, never a reset budget."""
+ timeout = float(timeout_seconds)
+ if not math.isfinite(timeout) or timeout <= 0:
+ raise ValueError("OAuth login timeout cap must be finite and positive.")
+ return OAuthLoginDeadline(
+ timeout,
+ timeout_message=self._timeout_message,
+ clock=self._clock,
+ sleeper=self._sleeper,
+ _expires_at=min(self._expires_at, self._clock() + timeout),
+ )
+
+
+def open_browser_async(url: str, opener: Callable[[str], object]) -> None:
+ """Launch a browser without letting OS integration consume the login budget."""
+
+ def _open() -> None:
+ try:
+ opener(url)
+ except Exception:
+ # The URL is always printed, so browser integration is best-effort.
+ return
+
+ threading.Thread(target=_open, daemon=True).start()
diff --git a/mobilerun/agent/utils/oauth/openai_oauth_llm.py b/mobilerun/agent/utils/oauth/openai_oauth_llm.py
index 70616b1d..c04f5edc 100644
--- a/mobilerun/agent/utils/oauth/openai_oauth_llm.py
+++ b/mobilerun/agent/utils/oauth/openai_oauth_llm.py
@@ -24,7 +24,8 @@
import time
import webbrowser
from dataclasses import dataclass
-from http.server import BaseHTTPRequestHandler, HTTPServer
+from http.server import BaseHTTPRequestHandler
+from http.server import ThreadingHTTPServer as HTTPServer
from pathlib import Path
from typing import Any, Dict, Optional
from urllib.parse import parse_qs, urlencode, urlparse
@@ -41,6 +42,11 @@
from llama_index.llms.openai.utils import to_openai_message_dicts
from mobilerun.agent.providers.registry import normalize_model_id_for_variant
+from mobilerun.agent.utils.oauth.login_timeout import (
+ OAuthLoginDeadline,
+ open_browser_async,
+)
+from mobilerun.config_manager.auth_profile_store import AuthProfileStore
from mobilerun.config_manager.credential_paths import OPENAI_OAUTH_CREDENTIAL_PATH
DEFAULT_OPENAI_OAUTH_ISSUER = "https://auth.openai.com"
@@ -56,6 +62,7 @@
DEFAULT_OPENAI_OAUTH_SCOPE = (
"openid profile email offline_access api.connectors.read api.connectors.invoke"
)
+_OPENAI_LOGIN_TIMEOUT_MESSAGE = "OpenAI OAuth login timed out."
def _b64_no_pad(raw: bytes) -> str:
@@ -80,7 +87,12 @@ def _is_headless_environment() -> bool:
return False
-def _tls_preflight(issuer: str, timeout: float = 5.0) -> None:
+def _tls_preflight(
+ issuer: str,
+ timeout: float = 5.0,
+ *,
+ deadline: OAuthLoginDeadline | None = None,
+) -> None:
"""Probe the OAuth issuer to detect TLS/certificate issues before login.
Raises RuntimeError on TLS certificate errors (with fix suggestions).
@@ -88,8 +100,13 @@ def _tls_preflight(issuer: str, timeout: float = 5.0) -> None:
"""
probe_url = f"{issuer.rstrip('/')}/oauth/authorize"
try:
- httpx.head(probe_url, follow_redirects=False, timeout=timeout)
+ request_timeout = deadline.remaining(cap=timeout) if deadline else timeout
+ httpx.head(probe_url, follow_redirects=False, timeout=request_timeout)
+ if deadline:
+ deadline.check()
except httpx.ConnectError as exc:
+ if deadline:
+ deadline.check()
err_str = str(exc).lower()
tls_indicators = (
"certificate",
@@ -116,11 +133,15 @@ def _tls_preflight(issuer: str, timeout: float = 5.0) -> None:
"The login flow may fail if there is a DNS or firewall issue."
)
except httpx.TimeoutException:
+ if deadline:
+ deadline.check()
print(
f"Warning: Connection to {probe_url} timed out.\n"
"The login flow may fail if there is a network issue."
)
except Exception as exc:
+ if deadline:
+ deadline.check()
# Unexpected error — warn but don't block.
print(f"Warning: TLS preflight check encountered an error: {exc}")
@@ -134,29 +155,47 @@ def _request_device_code(
http_client: Optional[httpx.Client] = None,
request_timeout: float = 15.0,
retries: int = 2,
+ *,
+ deadline: OAuthLoginDeadline | None = None,
) -> dict:
url = f"{issuer.rstrip('/')}/api/accounts/deviceauth/usercode"
post = http_client.post if http_client is not None else httpx.post
for attempt in range(1 + retries):
try:
+ timeout = (
+ deadline.remaining(cap=request_timeout) if deadline else request_timeout
+ )
response = post(
url,
headers={"Content-Type": "application/json"},
json={"client_id": client_id},
- timeout=request_timeout,
+ timeout=timeout,
)
+ if deadline:
+ deadline.check()
except (httpx.ConnectError, httpx.TimeoutException):
+ if deadline:
+ deadline.check()
if attempt < retries:
- time.sleep(2)
+ if deadline:
+ deadline.sleep(2)
+ else:
+ time.sleep(2)
continue
raise
if response.status_code == 404:
raise RuntimeError("Device code login is not enabled for this server.")
if response.status_code >= 500 and attempt < retries:
- time.sleep(2)
+ if deadline:
+ deadline.sleep(2)
+ else:
+ time.sleep(2)
continue
response.raise_for_status()
- return response.json()
+ payload = response.json()
+ if deadline:
+ deadline.check()
+ return payload
raise RuntimeError("Device code request failed after retries.")
@@ -168,55 +207,59 @@ def _poll_device_code(
http_client: Optional[httpx.Client] = None,
request_timeout: float = 15.0,
timeout_seconds: float = _DEVICE_CODE_TIMEOUT,
+ *,
+ deadline: OAuthLoginDeadline | None = None,
) -> dict:
url = f"{issuer.rstrip('/')}/api/accounts/deviceauth/token"
- effective_timeout = min(timeout_seconds, _DEVICE_CODE_TIMEOUT)
- deadline = time.time() + effective_timeout
+ login_deadline = deadline or OAuthLoginDeadline(
+ min(timeout_seconds, _DEVICE_CODE_TIMEOUT),
+ timeout_message=_OPENAI_LOGIN_TIMEOUT_MESSAGE,
+ )
post = http_client.post if http_client is not None else httpx.post
last_error: Optional[str] = None
- while time.time() < deadline:
- try:
- response = post(
- url,
- headers={"Content-Type": "application/json"},
- json={"device_auth_id": device_auth_id, "user_code": user_code},
- timeout=request_timeout,
- )
- except (httpx.ConnectError, httpx.TimeoutException) as exc:
- last_error = str(exc)
- remaining = deadline - time.time()
- if remaining <= 0:
- break
- time.sleep(min(interval, remaining))
- continue
- if response.status_code == 200:
- return response.json()
- # OpenAI returns 403/404 while the user hasn't completed browser auth.
- # This differs from RFC 8628's 400 + authorization_pending body, but
- # matches the observed behaviour of auth.openai.com/api/accounts/deviceauth/token.
- # TODO: handle slow_down (RFC 8628 §3.5) by increasing interval.
- if response.status_code in (403, 404):
- remaining = deadline - time.time()
- if remaining <= 0:
- break
- time.sleep(min(interval, remaining))
- continue
- if response.status_code >= 500:
- last_error = f"HTTP {response.status_code}"
- remaining = deadline - time.time()
- if remaining <= 0:
- break
- time.sleep(min(interval, remaining))
- continue
- response.raise_for_status()
-
- minutes = int(effective_timeout // 60)
- msg = f"Device code login timed out ({minutes} minutes)."
- if last_error:
- msg += f" Last response: {last_error}."
- raise TimeoutError(msg)
+ try:
+ while True:
+ try:
+ response = post(
+ url,
+ headers={"Content-Type": "application/json"},
+ json={
+ "device_auth_id": device_auth_id,
+ "user_code": user_code,
+ },
+ timeout=login_deadline.remaining(cap=request_timeout),
+ )
+ login_deadline.check()
+ except (httpx.ConnectError, httpx.TimeoutException) as exc:
+ login_deadline.check()
+ last_error = str(exc)
+ login_deadline.sleep(interval)
+ continue
+ if response.status_code == 200:
+ payload = response.json()
+ login_deadline.check()
+ return payload
+ # OpenAI returns 403/404 while the user hasn't completed browser auth.
+ # This differs from RFC 8628's 400 + authorization_pending body, but
+ # matches the observed behaviour of
+ # auth.openai.com/api/accounts/deviceauth/token.
+ # TODO: handle slow_down (RFC 8628 §3.5) by increasing interval.
+ if response.status_code in (403, 404):
+ login_deadline.sleep(interval)
+ continue
+ if response.status_code >= 500:
+ last_error = f"HTTP {response.status_code}"
+ login_deadline.sleep(interval)
+ continue
+ response.raise_for_status()
+ except TimeoutError:
+ if last_error:
+ raise TimeoutError(
+ f"{_OPENAI_LOGIN_TIMEOUT_MESSAGE} Last response: {last_error}."
+ ) from None
+ raise
@dataclass
@@ -273,6 +316,7 @@ def is_valid(self, skew_ms: int = 60_000) -> bool:
class OpenAIOAuthCredentialStore:
def __init__(self, path: str | Path = DEFAULT_OPENAI_OAUTH_CREDENTIAL_PATH) -> None:
self.path = Path(path).expanduser()
+ self._store = AuthProfileStore(self.path)
_NESTED_KEY = "openaiOauth"
@@ -280,7 +324,7 @@ def load(self) -> Optional[OpenAIOAuthCredentials]:
if not self.path.exists():
return None
- raw = json.loads(self.path.read_text())
+ raw = self._store.read_profile()
nested = raw.get(self._NESTED_KEY)
payload = nested if isinstance(nested, dict) else raw
try:
@@ -288,30 +332,18 @@ def load(self) -> Optional[OpenAIOAuthCredentials]:
except ValueError:
return None
- def save(self, credentials: OpenAIOAuthCredentials) -> None:
- self.path.parent.mkdir(parents=True, exist_ok=True)
-
- existing: dict = {}
- if self.path.exists():
- try:
- loaded = json.loads(self.path.read_text())
- if isinstance(loaded, dict):
- existing = loaded
- except Exception:
- existing = {}
-
- existing[self._NESTED_KEY] = credentials.to_dict()
-
- tmp_path = self.path.with_suffix(f"{self.path.suffix}.tmp")
- with open(tmp_path, "w", encoding="utf-8") as f:
- f.write(json.dumps(existing, indent=2))
- f.flush()
- os.fsync(f.fileno())
- os.replace(tmp_path, self.path)
- try:
- os.chmod(self.path, 0o600)
- except OSError:
- pass
+ def save(
+ self,
+ credentials: OpenAIOAuthCredentials,
+ *,
+ deadline: OAuthLoginDeadline | None = None,
+ ) -> None:
+ self._store.update_slot(
+ self._NESTED_KEY,
+ credentials.to_dict(),
+ lock_timeout=deadline.remaining() if deadline is not None else None,
+ before_commit=deadline.check if deadline is not None else None,
+ )
class OpenAIOAuthSessionManager:
@@ -385,10 +417,27 @@ def _compute_expiry_ms(cls, token: str, expires_in: Optional[Any]) -> Optional[i
return int(exp * 1000)
return None
- def set_initial_credentials(self, credentials: OpenAIOAuthCredentials) -> None:
- with self._lock:
+ def set_initial_credentials(
+ self,
+ credentials: OpenAIOAuthCredentials,
+ *,
+ deadline: OAuthLoginDeadline | None = None,
+ ) -> None:
+ if deadline is None:
+ with self._lock:
+ self._credentials = credentials
+ self.credential_store.save(credentials)
+ return
+
+ deadline.check()
+ if not self._lock.acquire(timeout=deadline.remaining()):
+ raise TimeoutError(_OPENAI_LOGIN_TIMEOUT_MESSAGE)
+ try:
+ deadline.check()
+ self.credential_store.save(credentials, deadline=deadline)
self._credentials = credentials
- self.credential_store.save(credentials)
+ finally:
+ self._lock.release()
def _load_cached_credentials(self) -> Optional[OpenAIOAuthCredentials]:
if self._credentials is not None:
@@ -469,6 +518,7 @@ def exchange_authorization_code(
code: str,
redirect_uri: str,
code_verifier: str,
+ deadline: OAuthLoginDeadline | None = None,
) -> OpenAIOAuthCredentials:
data = {
"grant_type": "authorization_code",
@@ -478,23 +528,32 @@ def exchange_authorization_code(
"code_verifier": code_verifier,
}
+ request_timeout = (
+ deadline.remaining(cap=self.request_timeout)
+ if deadline
+ else self.request_timeout
+ )
if self.http_client is not None:
response = self.http_client.post(
f"{self.issuer}/oauth/token",
headers={"Content-Type": "application/x-www-form-urlencoded"},
data=data,
- timeout=self.request_timeout,
+ timeout=request_timeout,
)
else:
response = httpx.post(
f"{self.issuer}/oauth/token",
headers={"Content-Type": "application/x-www-form-urlencoded"},
data=data,
- timeout=self.request_timeout,
+ timeout=request_timeout,
)
+ if deadline:
+ deadline.check()
response.raise_for_status()
payload = response.json()
+ if deadline:
+ deadline.check()
access = payload.get("access_token")
if not isinstance(access, str) or not access:
@@ -513,7 +572,7 @@ def exchange_authorization_code(
else access
),
)
- self.set_initial_credentials(credentials)
+ self.set_initial_credentials(credentials, deadline=deadline)
return credentials
@@ -661,19 +720,30 @@ def login(
callback_path: str = DEFAULT_OPENAI_OAUTH_CALLBACK_PATH,
redirect_host: str = DEFAULT_OPENAI_OAUTH_CALLBACK_HOST,
scope: str = DEFAULT_OPENAI_OAUTH_SCOPE,
+ deadline: OAuthLoginDeadline | None = None,
) -> OpenAIOAuthCredentials:
- _tls_preflight(self._oauth_manager.issuer)
+ login_deadline = deadline or OAuthLoginDeadline(
+ timeout_seconds,
+ timeout_message=_OPENAI_LOGIN_TIMEOUT_MESSAGE,
+ )
+ login_deadline.check()
+ _tls_preflight(self._oauth_manager.issuer, deadline=login_deadline)
# Headless environments: use device code flow (no local server needed)
use_device_code = _is_headless_environment() or os.environ.get(
"DROIDRUN_OAUTH_MANUAL", ""
).lower() in ("1", "true", "yes")
if use_device_code:
- return self._login_device_code(timeout_seconds=timeout_seconds)
+ return self._login_device_code(
+ open_browser=open_browser,
+ timeout_seconds=timeout_seconds,
+ deadline=login_deadline,
+ )
# Desktop: browser callback server
result: Dict[str, Optional[str]] = {"code": None, "state": None, "error": None}
done = threading.Event()
+ callback_lock = threading.Lock()
code_verifier, code_challenge = _pkce_pair()
state = _b64_no_pad(secrets.token_bytes(32))
@@ -685,12 +755,17 @@ def do_GET(self) -> None: # noqa: N802
self.end_headers()
return
- params = parse_qs(parsed.query)
- result["code"] = params.get("code", [None])[0]
- result["state"] = params.get("state", [None])[0]
- result["error"] = params.get("error", [None])[0]
-
- ok = result["code"] is not None and result["error"] is None
+ with callback_lock:
+ if done.is_set():
+ self.send_response(409)
+ self.end_headers()
+ return
+ params = parse_qs(parsed.query)
+ result["code"] = params.get("code", [None])[0]
+ result["state"] = params.get("state", [None])[0]
+ result["error"] = params.get("error", [None])[0]
+ ok = result["code"] is not None and result["error"] is None
+ done.set()
self.send_response(200 if ok else 400)
self.send_header("Content-Type", "text/html; charset=utf-8")
self.end_headers()
@@ -702,7 +777,6 @@ def do_GET(self) -> None: # noqa: N802
self.wfile.write(
b"Login failed. Return to your terminal.
"
)
- done.set()
def log_message(self, format: str, *args: Any) -> None: # noqa: A003
return
@@ -710,11 +784,17 @@ def log_message(self, format: str, *args: Any) -> None: # noqa: A003
try:
httpd = HTTPServer((callback_host, callback_port), _OAuthHandler)
except OSError as exc:
+ login_deadline.check()
print(
f"Could not bind callback server on {callback_host}:{callback_port} ({exc}). "
"Falling back to device code login."
)
- return self._login_device_code(timeout_seconds=timeout_seconds)
+ return self._login_device_code(
+ open_browser=open_browser,
+ timeout_seconds=timeout_seconds,
+ deadline=login_deadline,
+ )
+ login_deadline.check()
actual_port = httpd.server_address[1]
redirect_uri = f"http://{redirect_host}:{actual_port}{callback_path}"
@@ -728,17 +808,19 @@ def log_message(self, format: str, *args: Any) -> None: # noqa: A003
)
server_thread = threading.Thread(target=httpd.serve_forever, daemon=True)
+ httpd.daemon_threads = True
server_thread.start()
try:
print(f"Open this URL to login:\n{auth_url}\n")
if open_browser:
- webbrowser.open(auth_url)
+ open_browser_async(auth_url, webbrowser.open)
- if not done.wait(timeout=timeout_seconds):
+ if not done.wait(timeout=login_deadline.remaining()):
raise TimeoutError(
- "OAuth login timed out before callback was received."
+ "OpenAI OAuth login timed out before callback was received."
)
+ login_deadline.check()
if result["error"]:
raise RuntimeError(f"OAuth callback returned error: {result['error']}")
@@ -753,6 +835,7 @@ def log_message(self, format: str, *args: Any) -> None: # noqa: A003
code=result["code"],
redirect_uri=redirect_uri,
code_verifier=code_verifier,
+ deadline=login_deadline,
)
if creds.account_id:
object.__setattr__(self, "_oauth_account_id", creds.account_id)
@@ -764,7 +847,9 @@ def log_message(self, format: str, *args: Any) -> None: # noqa: A003
def _login_device_code(
self,
*,
+ open_browser: bool = False,
timeout_seconds: float = _DEVICE_CODE_TIMEOUT,
+ deadline: OAuthLoginDeadline | None = None,
) -> OpenAIOAuthCredentials:
"""Device Code login for headless/SSH environments.
@@ -772,12 +857,18 @@ def _login_device_code(
"""
mgr = self._oauth_manager
http_client = mgr.http_client
+ login_deadline = deadline or OAuthLoginDeadline(
+ timeout_seconds,
+ timeout_message=_OPENAI_LOGIN_TIMEOUT_MESSAGE,
+ )
+ login_deadline.check()
device_resp = _request_device_code(
mgr.issuer,
mgr.client_id,
http_client=http_client,
request_timeout=mgr.request_timeout,
+ deadline=login_deadline,
)
device_auth_id = device_resp.get("device_auth_id")
if not device_auth_id:
@@ -787,14 +878,19 @@ def _login_device_code(
if not user_code:
raise RuntimeError("Device code response missing 'user_code'.")
try:
- interval = int(str(device_resp.get("interval", "5")).strip())
+ interval = max(1, int(str(device_resp.get("interval", "5")).strip()))
except (TypeError, ValueError):
interval = 5
try:
server_expires = int(device_resp["expires_in"])
except (KeyError, TypeError, ValueError):
server_expires = _DEVICE_CODE_TIMEOUT
- effective_timeout = min(timeout_seconds, _DEVICE_CODE_TIMEOUT, server_expires)
+ if server_expires <= 0:
+ raise TimeoutError("OpenAI OAuth device authorization expired.")
+ device_deadline = login_deadline.limited_to(
+ min(_DEVICE_CODE_TIMEOUT, server_expires)
+ )
+ effective_timeout = device_deadline.remaining()
verification_url = (
device_resp.get("verification_uri")
or device_resp.get("verification_url")
@@ -813,6 +909,8 @@ def _login_device_code(
f"\n2. Enter this code (expires in {expires_str}):\n {user_code}\n"
f"\nDevice codes are a common phishing target. Never share this code.\n"
)
+ if open_browser:
+ open_browser_async(str(verification_url), webbrowser.open)
token_resp = _poll_device_code(
mgr.issuer,
@@ -821,7 +919,7 @@ def _login_device_code(
interval,
http_client=http_client,
request_timeout=mgr.request_timeout,
- timeout_seconds=effective_timeout,
+ deadline=device_deadline,
)
auth_code = token_resp.get("authorization_code")
@@ -837,6 +935,7 @@ def _login_device_code(
code=auth_code,
redirect_uri=redirect_uri,
code_verifier=code_verifier,
+ deadline=device_deadline,
)
if creds.account_id:
object.__setattr__(self, "_oauth_account_id", creds.account_id)
diff --git a/mobilerun/cli/configure_wizard.py b/mobilerun/cli/configure_wizard.py
index a6bd8f73..115cc9b5 100644
--- a/mobilerun/cli/configure_wizard.py
+++ b/mobilerun/cli/configure_wizard.py
@@ -49,6 +49,7 @@ class ConfigureWizardCallbacks:
run_openai_oauth_login: Callable[..., None]
run_anthropic_oauth_login: Callable[..., None]
run_gemini_oauth_login: Callable[..., None]
+ run_grok_oauth_login: Callable[..., None] | None = None
@dataclass
@@ -107,9 +108,10 @@ def _print_configure_summary(
used_advanced_settings: bool,
) -> None:
advanced_line = "Yes" if used_advanced_settings else "No"
+ provider_detail = "" if provider_label == "XAI" else f" ({variant_id})"
console.print(
Panel(
- f"Provider: {provider_label} ({variant_id})\n"
+ f"Provider: {provider_label}{provider_detail}\n"
f"Model: {model}\n"
f"Advanced settings changed: {advanced_line}",
title="Configuration Saved",
@@ -242,6 +244,7 @@ def _prompt_base_url_for_variant(variant: Any) -> str:
"gemini_oauth_code_assist": "geminiAntigravityOauth",
"openai_oauth": "openaiOauth",
"anthropic_oauth": "claudeAiOauth",
+ "xai_oauth": "grokOauth",
}
@@ -321,6 +324,13 @@ def _prepare_variant_auth(
callbacks.run_gemini_oauth_login(
credential_path=credential_path, model=selected_model
)
+ elif variant.id == "xai_oauth" and credential_path:
+ if callbacks.run_grok_oauth_login is None:
+ raise RuntimeError("XAI OAuth login callback is not configured.")
+ callbacks.run_grok_oauth_login(
+ credential_path=credential_path,
+ model=selected_model,
+ )
def _set_profile_max_tokens(profile: Any, value: int) -> None:
@@ -547,6 +557,8 @@ def _configure_provider_model(
)
if state.selected_model == _BACK:
state.selected_model = None
+ if provider_is_fixed and (auth_mode_is_fixed or len(modes) == 1):
+ return False
if auth_mode_is_fixed or len(modes) == 1:
if not provider_is_fixed:
state.family_id = None
@@ -682,10 +694,11 @@ def run_configure_wizard(
if model_is_fixed:
state.selected_model = model
- # When CLI flags fully specify provider+model, run the flow automatically
- # and save without showing the menu.
+ # Enter a partially fixed provider flow only when another flag establishes
+ # where that flow should resume. A provider-only invocation starts at the
+ # top-level menu, matching the fully interactive wizard.
provider_configured = False
- if provider_is_fixed and model_is_fixed:
+ if provider_is_fixed and (auth_mode_is_fixed or model_is_fixed):
provider_configured = _configure_provider_model(
console,
config,
diff --git a/mobilerun/cli/main.py b/mobilerun/cli/main.py
index 419023da..9a2f13d4 100644
--- a/mobilerun/cli/main.py
+++ b/mobilerun/cli/main.py
@@ -5,6 +5,7 @@
import asyncio
import importlib.metadata
import logging
+import math
import os
import sys
import tomllib
@@ -53,6 +54,7 @@
from mobilerun.cli.oauth_actions import (
run_anthropic_setup_token_oauth,
run_gemini_oauth_login,
+ run_grok_oauth_login,
run_openai_oauth_login,
save_anthropic_setup_token,
)
@@ -60,6 +62,7 @@
from mobilerun.config_manager.credential_paths import (
ANTHROPIC_OAUTH_CREDENTIAL_PATH,
GEMINI_OAUTH_CREDENTIAL_PATH,
+ GROK_OAUTH_CREDENTIAL_PATH,
)
from mobilerun.log_handlers import CLILogHandler, configure_logging
from mobilerun.macro.cli import macro_cli
@@ -75,6 +78,16 @@
console = Console()
+def _validate_oauth_timeout(
+ _ctx: click.Context,
+ _param: click.Parameter,
+ value: float,
+) -> float:
+ if not math.isfinite(value) or value <= 0:
+ raise click.BadParameter("must be a finite number greater than zero")
+ return value
+
+
def _force_screenshot_only_vision(config: MobileConfig) -> None:
config.agent.vision_only = True
config.agent.manager.vision = True
@@ -416,12 +429,25 @@ def _run_gemini_oauth_login(credential_path: str, model: str | None, **kwargs) -
def _run_anthropic_oauth_login(credential_path: str, **kwargs) -> None:
"""Run the full Anthropic OAuth flow inline and save the token."""
- console.print("[blue]Opening browser for Anthropic login...[/]")
- token = run_anthropic_setup_token_oauth(**kwargs)
- save_anthropic_setup_token(credential_path, token)
+ console.print("[blue]Starting Anthropic login...[/]")
+ run_anthropic_setup_token_oauth(
+ credential_path=credential_path,
+ **kwargs,
+ )
_print_oauth_login_success("Anthropic", credential_path)
+def _run_grok_oauth_login(
+ credential_path: str, model: str | None = None, **kwargs
+) -> None:
+ run_grok_oauth_login(
+ credential_path=credential_path,
+ model=model,
+ **kwargs,
+ )
+ _print_oauth_login_success("XAI", credential_path)
+
+
try:
_available_agents = list_agents()
except Exception:
@@ -447,7 +473,7 @@ def _run_anthropic_oauth_login(credential_path: str, **kwargs) -> None:
@click.option(
"--provider",
"-p",
- help="LLM provider (OpenAI, openai_oauth, Ollama, Anthropic, anthropic_oauth, GoogleGenAI, gemini_oauth_code_assist, DeepSeek)",
+ help="LLM provider (OpenAI, openai_oauth, XAI, Ollama, Anthropic, anthropic_oauth, GoogleGenAI, gemini_oauth_code_assist, DeepSeek)",
default=None,
)
@click.option(
@@ -1012,7 +1038,7 @@ async def doctor(device: str | None, debug: bool | None):
"--provider",
type=str,
default=None,
- help="Provider family (gemini, openai, anthropic, ollama, openai_like, minimax, zai).",
+ help="Provider family (gemini, openai, anthropic, XAI, ollama, openai_like, minimax, zai).",
)
@click.option(
"--auth-mode",
@@ -1052,6 +1078,7 @@ def configure(
run_openai_oauth_login=_run_openai_oauth_login,
run_anthropic_oauth_login=_run_anthropic_oauth_login,
run_gemini_oauth_login=_run_gemini_oauth_login,
+ run_grok_oauth_login=_run_grok_oauth_login,
),
provider=provider,
auth_mode=auth_mode,
@@ -1074,9 +1101,10 @@ def configure(
@click.option(
"--timeout",
type=float,
+ callback=_validate_oauth_timeout,
default=300.0,
show_default=True,
- help="Max seconds to wait for the browser callback.",
+ help="Max seconds allowed for the complete OpenAI OAuth login.",
)
@click.option(
"--callback-host",
@@ -1136,13 +1164,36 @@ def configure_openai(
default=None,
help="Anthropic setup-token value. If provided, skips the OAuth flow.",
)
-def configure_anthropic(credential_path: str, token: str | None):
+@click.option(
+ "--timeout",
+ type=float,
+ callback=_validate_oauth_timeout,
+ default=300.0,
+ show_default=True,
+ help="Max seconds allowed for the complete Anthropic OAuth login.",
+)
+@click.option(
+ "--open-browser/--no-browser",
+ default=True,
+ show_default=True,
+ help="Open the Anthropic authorization URL automatically.",
+)
+def configure_anthropic(
+ credential_path: str,
+ token: str | None,
+ timeout: float,
+ open_browser: bool,
+):
"""Log in to Anthropic via OAuth (or pass --token to save a setup-token)."""
if token:
save_anthropic_setup_token(credential_path, token)
_print_oauth_login_success("Anthropic", credential_path)
else:
- _run_anthropic_oauth_login(credential_path=credential_path)
+ _run_anthropic_oauth_login(
+ credential_path=credential_path,
+ timeout=timeout,
+ open_browser=open_browser,
+ )
@configure.command("gemini")
@@ -1158,9 +1209,10 @@ def configure_anthropic(credential_path: str, token: str | None):
@click.option(
"--timeout",
type=float,
+ callback=_validate_oauth_timeout,
default=300.0,
show_default=True,
- help="Max seconds to wait for the browser callback.",
+ help="Max seconds allowed for the complete Gemini OAuth login.",
)
@click.option(
"--callback-host",
@@ -1208,6 +1260,53 @@ def configure_gemini(
)
+@configure.command("xai")
+@click.option(
+ "--credential-path",
+ default=str(GROK_OAUTH_CREDENTIAL_PATH),
+ show_default=True,
+ help="Where to store XAI OAuth credentials.",
+)
+@click.option(
+ "--model", default=None, help="Optional model override for later API calls."
+)
+@click.option(
+ "--timeout",
+ type=float,
+ callback=_validate_oauth_timeout,
+ default=300.0,
+ show_default=True,
+ help="Max seconds allowed for the complete xAI OAuth login.",
+)
+@click.option(
+ "--open-browser/--no-browser",
+ default=True,
+ show_default=True,
+ help="Open the xAI authorization URL automatically.",
+)
+@click.option(
+ "--device-code",
+ is_flag=True,
+ default=False,
+ help="Use xAI's device-code flow for SSH or other headless environments.",
+)
+def configure_xai(
+ credential_path: str,
+ model: str | None,
+ timeout: float,
+ open_browser: bool,
+ device_code: bool,
+):
+ """Log in to XAI with Mobilerun OAuth."""
+ _run_grok_oauth_login(
+ credential_path=credential_path,
+ model=model,
+ timeout=timeout,
+ open_browser=open_browser,
+ device_code=device_code,
+ )
+
+
async def test(
command: str,
config_path: str | None = None,
diff --git a/mobilerun/cli/oauth_actions.py b/mobilerun/cli/oauth_actions.py
index 91105cde..2ffd1144 100644
--- a/mobilerun/cli/oauth_actions.py
+++ b/mobilerun/cli/oauth_actions.py
@@ -1,9 +1,5 @@
from __future__ import annotations
-import json
-import os
-from pathlib import Path
-
from mobilerun.agent.utils.oauth.anthropic_oauth_llm import (
DEFAULT_SETUP_TOKEN_SCOPE,
AnthropicOAuthLLM,
@@ -14,6 +10,11 @@
from mobilerun.agent.utils.oauth.gemini_oauth_code_assist_llm import (
GeminiOAuthCodeAssistLLM,
)
+from mobilerun.agent.utils.oauth.grok_oauth_llm import (
+ DEFAULT_GROK_MODEL,
+ GrokOAuth,
+)
+from mobilerun.agent.utils.oauth.login_timeout import OAuthLoginDeadline
from mobilerun.agent.utils.oauth.openai_oauth_llm import (
DEFAULT_OPENAI_OAUTH_CALLBACK_HOST,
DEFAULT_OPENAI_OAUTH_CALLBACK_PATH,
@@ -21,6 +22,7 @@
DEFAULT_OPENAI_OAUTH_CREDENTIAL_PATH,
OpenAIOAuth,
)
+from mobilerun.config_manager.auth_profile_store import AuthProfileStore
SETUP_TOKEN_EXPIRES_IN_SECONDS = 365 * 24 * 60 * 60
@@ -34,14 +36,23 @@ def run_openai_oauth_login(
callback_path: str = DEFAULT_OPENAI_OAUTH_CALLBACK_PATH,
open_browser: bool = True,
) -> None:
- llm = OpenAIOAuth(model=model, oauth_credential_path=credential_path)
+ deadline = OAuthLoginDeadline(
+ timeout,
+ timeout_message="OpenAI OAuth login timed out.",
+ )
+ llm = OpenAIOAuth(
+ model=model,
+ oauth_credential_path=credential_path,
+ timeout=timeout,
+ )
llm.login(
open_browser=open_browser,
- timeout_seconds=timeout,
+ timeout_seconds=deadline.remaining(),
callback_host=callback_host,
callback_port=callback_port,
callback_path=callback_path,
redirect_host=callback_host,
+ deadline=deadline,
)
@@ -54,80 +65,134 @@ def run_gemini_oauth_login(
callback_path: str = "/oauth2callback",
open_browser: bool = True,
) -> None:
+ deadline = OAuthLoginDeadline(
+ timeout,
+ timeout_message="Gemini OAuth login timed out.",
+ )
llm = GeminiOAuthCodeAssistLLM(
model=model or GEMINI_OAUTH_DEFAULT_MODEL,
credential_path=credential_path,
+ timeout=timeout,
)
- llm.login(
+ access_token = llm.login(
open_browser=open_browser,
- timeout_seconds=timeout,
+ timeout_seconds=deadline.remaining(),
callback_host=callback_host,
callback_port=callback_port,
callback_path=callback_path,
+ deadline=deadline,
+ persist_credentials=False,
)
# Verify the Antigravity consumer entitlement resolves before declaring
# success (catches scope / header / endpoint problems at login time). Raise
# on failure so the caller does not print a misleading success message.
try:
- models = llm.fetch_available_models()
+ models = llm.fetch_available_models(
+ deadline=deadline,
+ access_token=access_token,
+ )
+ except TimeoutError:
+ raise
except Exception as exc: # noqa: BLE001
raise RuntimeError(
- "Gemini OAuth login saved a token, but the Antigravity entitlement "
+ "Gemini OAuth login received a token, but the Antigravity entitlement "
f"check (fetchAvailableModels) failed: {exc}. The login is not "
"usable; verify your Google One / AI access and retry."
) from exc
+ if not models:
+ raise RuntimeError(
+ "Gemini OAuth login returned no usable models; credentials were not saved."
+ )
+ llm._persist_credentials(deadline=deadline)
print(f"✓ Gemini (Antigravity) login OK — {len(models)} models available.")
+def run_grok_oauth_login(
+ credential_path: str,
+ model: str | None,
+ timeout: float = 300.0,
+ open_browser: bool = True,
+ device_code: bool = False,
+ no_browser: bool | None = None,
+) -> None:
+ """Authenticate directly with xAI and save Mobilerun-owned credentials."""
+ deadline = OAuthLoginDeadline(
+ timeout,
+ timeout_message="XAI OAuth login timed out.",
+ )
+ llm = GrokOAuth(
+ model=model or DEFAULT_GROK_MODEL,
+ oauth_credential_path=credential_path,
+ timeout=timeout,
+ )
+ llm.login(
+ open_browser=(
+ open_browser if no_browser is None else open_browser and not no_browser
+ ),
+ timeout_seconds=deadline.remaining(),
+ device_code=device_code,
+ deadline=deadline,
+ )
+
+
def run_anthropic_setup_token_oauth(
*,
+ credential_path: str | None = None,
timeout: float = 300.0,
callback_host: str = "127.0.0.1",
callback_port: int = 0,
callback_path: str = "/callback",
open_browser: bool = True,
) -> str:
+ deadline = OAuthLoginDeadline(
+ timeout,
+ timeout_message="Anthropic OAuth login timed out.",
+ )
llm = AnthropicOAuthLLM(
credential_path=None,
authorize_url="https://claude.com/cai/oauth/authorize",
login_scope=DEFAULT_SETUP_TOKEN_SCOPE,
+ timeout=timeout,
)
- return llm.login(
+ token = llm.login(
open_browser=open_browser,
- timeout_seconds=timeout,
+ timeout_seconds=deadline.remaining(),
callback_host=callback_host,
callback_port=callback_port,
callback_path=callback_path,
expires_in=SETUP_TOKEN_EXPIRES_IN_SECONDS,
+ deadline=deadline,
)
+ if credential_path is not None:
+ save_anthropic_setup_token(
+ credential_path,
+ token,
+ deadline=deadline,
+ )
+ return token
-def save_anthropic_setup_token(credential_path: str, token: str) -> None:
- cred_path = Path(credential_path).expanduser()
- cred_path.parent.mkdir(parents=True, exist_ok=True)
-
- existing: dict[str, object] = {}
- if cred_path.exists():
- try:
- loaded = json.loads(cred_path.read_text(encoding="utf-8"))
- if isinstance(loaded, dict):
- existing = loaded
- except Exception:
- existing = {}
-
- existing["claudeAiOauth"] = {
- "accessToken": token,
- "refreshToken": None,
- "expiresAt": None,
- "scopes": [],
- }
- cred_path.write_text(json.dumps(existing, indent=2), encoding="utf-8")
- os.chmod(cred_path, 0o600)
+def save_anthropic_setup_token(
+ credential_path: str,
+ token: str,
+ *,
+ deadline: OAuthLoginDeadline | None = None,
+) -> None:
+ AuthProfileStore(credential_path).update_slot(
+ "claudeAiOauth",
+ {
+ "accessToken": token,
+ "refreshToken": None,
+ "expiresAt": None,
+ "scopes": [],
+ },
+ lock_timeout=deadline.remaining() if deadline is not None else None,
+ before_commit=deadline.check if deadline is not None else None,
+ )
def run_anthropic_oauth_setup(credential_path: str) -> None:
- token = run_anthropic_setup_token_oauth()
- save_anthropic_setup_token(credential_path, token)
+ run_anthropic_setup_token_oauth(credential_path=credential_path)
def get_default_openai_credential_path() -> str:
diff --git a/mobilerun/config_manager/auth_profile_store.py b/mobilerun/config_manager/auth_profile_store.py
new file mode 100644
index 00000000..5b50d158
--- /dev/null
+++ b/mobilerun/config_manager/auth_profile_store.py
@@ -0,0 +1,202 @@
+"""Safe shared storage for Mobilerun authentication profiles.
+
+All OAuth providers and saved API keys share one JSON object. Writes therefore
+need to be serialized across processes and must preserve slots owned by other
+providers.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import tempfile
+from contextlib import AbstractContextManager
+from pathlib import Path
+from typing import Any, Callable, TypeVar
+
+from filelock import FileLock
+
+
+class AuthProfileFormatError(ValueError):
+ """Raised when an existing auth-profiles file is not a JSON object."""
+
+
+_T = TypeVar("_T")
+_FCHMOD: Callable[[int, int], None] | None = getattr(os, "fchmod", None)
+
+
+class AuthProfileTransaction(AbstractContextManager["AuthProfileTransaction"]):
+ """A locked read/modify/write transaction over an auth profile file."""
+
+ def __init__(
+ self,
+ store: "AuthProfileStore",
+ *,
+ lock_timeout: float | None = None,
+ before_commit: Callable[[], None] | None = None,
+ ) -> None:
+ self._store = store
+ self._lock = FileLock(str(store.lock_path))
+ self._lock_timeout = lock_timeout
+ self._before_commit = before_commit
+ self._profile: dict[str, Any] = {}
+ self._dirty = False
+
+ def __enter__(self) -> "AuthProfileTransaction":
+ self._store.path.parent.mkdir(parents=True, exist_ok=True)
+ if self._lock_timeout is None:
+ self._lock.acquire()
+ else:
+ self._lock.acquire(timeout=self._lock_timeout)
+ try:
+ # The lock file has no secrets, but keeping it private prevents
+ # other local users from deliberately interfering with writers.
+ try:
+ os.chmod(self._store.lock_path, 0o600)
+ except OSError:
+ pass
+ self._profile = self._store._read_unlocked()
+ except BaseException:
+ self._lock.release()
+ raise
+ return self
+
+ @property
+ def profile(self) -> dict[str, Any]:
+ """The locked profile object. Mutate it only via :meth:`update`."""
+ return self._profile
+
+ def get_slot(self, slot: str) -> dict[str, Any] | None:
+ value = self._profile.get(slot)
+ return dict(value) if isinstance(value, dict) else None
+
+ def set_slot(self, slot: str, payload: dict[str, Any]) -> None:
+ self._profile[slot] = dict(payload)
+ self._dirty = True
+
+ def update(self, updater: Callable[[dict[str, Any]], _T]) -> _T:
+ result = updater(self._profile)
+ self._dirty = True
+ return result
+
+ def __exit__(self, exc_type: object, exc: object, traceback: object) -> None:
+ try:
+ if exc_type is None and self._dirty:
+ self._store._write_unlocked(
+ self._profile,
+ before_commit=self._before_commit,
+ )
+ finally:
+ self._lock.release()
+ return None
+
+
+class AuthProfileStore:
+ """Cross-process-safe JSON object store used by all auth providers."""
+
+ def __init__(self, path: str | Path) -> None:
+ self.path = Path(path).expanduser()
+ self.lock_path = self.path.with_name(f"{self.path.name}.lock")
+
+ def transaction(
+ self,
+ *,
+ lock_timeout: float | None = None,
+ before_commit: Callable[[], None] | None = None,
+ ) -> AuthProfileTransaction:
+ return AuthProfileTransaction(
+ self,
+ lock_timeout=lock_timeout,
+ before_commit=before_commit,
+ )
+
+ def read_profile(self) -> dict[str, Any]:
+ with self.transaction() as transaction:
+ return dict(transaction.profile)
+
+ def read_slot(self, slot: str) -> dict[str, Any] | None:
+ with self.transaction() as transaction:
+ return transaction.get_slot(slot)
+
+ def update_slot(
+ self,
+ slot: str,
+ payload: dict[str, Any],
+ *,
+ lock_timeout: float | None = None,
+ before_commit: Callable[[], None] | None = None,
+ ) -> None:
+ with self.transaction(
+ lock_timeout=lock_timeout,
+ before_commit=before_commit,
+ ) as transaction:
+ transaction.set_slot(slot, payload)
+
+ def update_profile(self, updater: Callable[[dict[str, Any]], _T]) -> _T:
+ with self.transaction() as transaction:
+ return transaction.update(updater)
+
+ def _read_unlocked(self) -> dict[str, Any]:
+ if not self.path.exists():
+ return {}
+ try:
+ payload = json.loads(self.path.read_text(encoding="utf-8"))
+ except (OSError, json.JSONDecodeError) as exc:
+ raise AuthProfileFormatError(
+ f"Cannot update malformed authentication profile {self.path}. "
+ "Repair or remove the file, then retry."
+ ) from exc
+ if not isinstance(payload, dict):
+ raise AuthProfileFormatError(
+ f"Authentication profile {self.path} must contain a JSON object."
+ )
+ return payload
+
+ def _write_unlocked(
+ self,
+ profile: dict[str, Any],
+ *,
+ before_commit: Callable[[], None] | None = None,
+ ) -> None:
+ self.path.parent.mkdir(parents=True, exist_ok=True)
+ fd: int | None = None
+ tmp_path: Path | None = None
+ try:
+ fd, raw_tmp_path = tempfile.mkstemp(
+ dir=self.path.parent,
+ prefix=f".{self.path.name}.",
+ suffix=".tmp",
+ )
+ tmp_path = Path(raw_tmp_path)
+ if _FCHMOD is not None:
+ _FCHMOD(fd, 0o600)
+ with os.fdopen(fd, "w", encoding="utf-8") as handle:
+ fd = None
+ json.dump(profile, handle, indent=2)
+ handle.write("\n")
+ handle.flush()
+ os.fsync(handle.fileno())
+ if before_commit is not None:
+ before_commit()
+ os.replace(tmp_path, self.path)
+ tmp_path = None
+ os.chmod(self.path, 0o600)
+
+ # Persist the directory entry as well as the file contents.
+ try:
+ directory_fd = os.open(self.path.parent, os.O_RDONLY)
+ except OSError:
+ directory_fd = None
+ if directory_fd is not None:
+ try:
+ os.fsync(directory_fd)
+ finally:
+ os.close(directory_fd)
+ finally:
+ if fd is not None:
+ os.close(fd)
+ if tmp_path is not None:
+ try:
+ tmp_path.unlink()
+ except FileNotFoundError:
+ pass
diff --git a/mobilerun/config_manager/credential_paths.py b/mobilerun/config_manager/credential_paths.py
index 12163e13..8cf5d035 100644
--- a/mobilerun/config_manager/credential_paths.py
+++ b/mobilerun/config_manager/credential_paths.py
@@ -15,3 +15,4 @@
OPENAI_OAUTH_CREDENTIAL_PATH = AUTH_PROFILES_PATH
ANTHROPIC_OAUTH_CREDENTIAL_PATH = AUTH_PROFILES_PATH
GEMINI_OAUTH_CREDENTIAL_PATH = AUTH_PROFILES_PATH
+GROK_OAUTH_CREDENTIAL_PATH = AUTH_PROFILES_PATH
diff --git a/mobilerun/config_manager/env_keys.py b/mobilerun/config_manager/env_keys.py
index b89f6c42..96effc52 100644
--- a/mobilerun/config_manager/env_keys.py
+++ b/mobilerun/config_manager/env_keys.py
@@ -2,16 +2,17 @@
from __future__ import annotations
-import json
import os
from dataclasses import dataclass
+from mobilerun.config_manager.auth_profile_store import AuthProfileStore
from mobilerun.config_manager.credential_paths import AUTH_PROFILES_PATH
API_KEY_ENV_VARS = {
"google": "GOOGLE_API_KEY",
"gemini": "GEMINI_API_KEY",
"openai": "OPENAI_API_KEY",
+ "xai": "XAI_API_KEY",
"anthropic": "ANTHROPIC_API_KEY",
"zai": "ZAI_API_KEY",
"minimax": "MINIMAX_API_KEY",
@@ -33,7 +34,7 @@ def _load_saved_api_keys() -> dict[str, str]:
if not AUTH_PROFILES_PATH.exists():
return {}
try:
- data = json.loads(AUTH_PROFILES_PATH.read_text(encoding="utf-8"))
+ data = AuthProfileStore(AUTH_PROFILES_PATH).read_profile()
section = data.get(_API_KEYS_SECTION)
if isinstance(section, dict):
return {k: str(v) for k, v in section.items() if v}
@@ -96,38 +97,25 @@ def save_env_keys(keys: dict[str, str]) -> None:
Args:
keys: Dict mapping slot name (e.g. "google") to key value.
"""
- AUTH_PROFILES_PATH.parent.mkdir(parents=True, exist_ok=True)
-
- existing: dict = {}
- if AUTH_PROFILES_PATH.exists():
- try:
- loaded = json.loads(AUTH_PROFILES_PATH.read_text(encoding="utf-8"))
- if isinstance(loaded, dict):
- existing = loaded
- except Exception:
- existing = {}
-
- api_keys = existing.get(_API_KEYS_SECTION, {})
- if not isinstance(api_keys, dict):
- api_keys = {}
-
- for slot, val in keys.items():
- env_var = API_KEY_ENV_VARS.get(slot)
- if not env_var:
- continue
+ known_keys = {slot: val for slot, val in keys.items() if slot in API_KEY_ENV_VARS}
+
+ def _update(existing: dict) -> None:
+ api_keys = existing.get(_API_KEYS_SECTION, {})
+ if not isinstance(api_keys, dict):
+ api_keys = {}
+ for slot, val in known_keys.items():
+ if val:
+ api_keys[slot] = val
+ else:
+ api_keys.pop(slot, None)
+ existing[_API_KEYS_SECTION] = api_keys
+
+ # Persist first. In particular, a malformed file must not be replaced or
+ # leave the current process believing a key was saved successfully.
+ AuthProfileStore(AUTH_PROFILES_PATH).update_profile(_update)
+ for slot, val in known_keys.items():
+ env_var = API_KEY_ENV_VARS[slot]
if val:
- api_keys[slot] = val
os.environ[env_var] = val
else:
- api_keys.pop(slot, None)
os.environ.pop(env_var, None)
-
- existing[_API_KEYS_SECTION] = api_keys
-
- tmp_path = AUTH_PROFILES_PATH.with_suffix(".json.tmp")
- tmp_path.write_text(json.dumps(existing, indent=2), encoding="utf-8")
- os.replace(tmp_path, AUTH_PROFILES_PATH)
- try:
- os.chmod(AUTH_PROFILES_PATH, 0o600)
- except OSError:
- pass
diff --git a/pyproject.toml b/pyproject.toml
index 87ef8102..e3571fa6 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -12,6 +12,8 @@ dependencies = [
"InquirerPy>=0.3.4",
"arize-phoenix>=12.3.0",
"httpx>=0.27.0",
+ "filelock>=3.20.0",
+ "PyJWT[crypto]>=2.10.0",
"llama-index-callbacks-arize-phoenix>=0.6.1",
"llama-index-workflows>=2.16.0,<3.0.0",
"aiofiles>=25.1.0",
diff --git a/tests/test_grok_api.py b/tests/test_grok_api.py
new file mode 100644
index 00000000..6e946156
--- /dev/null
+++ b/tests/test_grok_api.py
@@ -0,0 +1,572 @@
+from __future__ import annotations
+
+import asyncio
+import json
+from types import SimpleNamespace
+from typing import Any
+
+import httpx
+import pytest
+from llama_index.core.base.llms.types import (
+ ChatMessage,
+ ImageBlock,
+ MessageRole,
+ TextBlock,
+)
+from openai.types.responses import Response
+from openai.types.responses.response_completed_event import ResponseCompletedEvent
+from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall
+from openai.types.responses.response_usage import (
+ InputTokensDetails,
+ OutputTokensDetails,
+ ResponseUsage,
+)
+
+from mobilerun.agent.providers.grok import XAI_API_BASE
+from mobilerun.agent.providers.registry import (
+ VARIANT_ENV_KEY_SLOT,
+ list_models_for_variant,
+ normalize_model_id_for_variant,
+ resolve_provider_variant,
+)
+from mobilerun.agent.providers.setup_service import (
+ SetupSelection,
+ create_profile_for_variant,
+)
+from mobilerun.agent.usage import get_usage_from_response
+from mobilerun.agent.utils.llm_picker import load_llm, normalize_provider_name
+from mobilerun.config_manager import env_keys
+
+
+def _xai_completed_response(*, usage: ResponseUsage) -> Response:
+ return Response(
+ id="response-id",
+ created_at=0,
+ error=None,
+ incomplete_details=None,
+ instructions=None,
+ metadata={},
+ model="grok-4.5",
+ object="response",
+ output=[
+ ResponseFunctionToolCall(
+ arguments='{"value":"ok"}',
+ call_id="call-id",
+ name="inspect",
+ type="function_call",
+ id="item-id",
+ status="completed",
+ )
+ ],
+ parallel_tool_calls=True,
+ temperature=0.4,
+ tool_choice="auto",
+ tools=[],
+ top_p=0.6,
+ status="completed",
+ usage=usage,
+ )
+
+
+def _xai_usage() -> ResponseUsage:
+ return ResponseUsage(
+ input_tokens=7,
+ input_tokens_details=InputTokensDetails(cached_tokens=0),
+ output_tokens=4,
+ output_tokens_details=OutputTokensDetails(reasoning_tokens=0),
+ total_tokens=11,
+ )
+
+
+def test_grok_api_key_variant_is_first_class_xai_responses_provider() -> None:
+ variant = resolve_provider_variant("xai", "api_key")
+
+ assert variant.id == "XAI"
+ assert variant.runtime_provider_name == "XAI"
+ assert variant.default_model == "grok-4.5"
+ assert list_models_for_variant("xai", "api_key") == ("grok-4.5",)
+ assert variant.requires_api_key is True
+ assert variant.base_url == XAI_API_BASE
+ assert VARIANT_ENV_KEY_SLOT[variant.id] == "xai"
+ assert env_keys.API_KEY_ENV_VARS["xai"] == "XAI_API_KEY"
+
+
+def test_grok_oauth_variant_shares_the_canonical_model_catalog() -> None:
+ variant = resolve_provider_variant("xai", "oauth")
+
+ assert variant.id == "xai_oauth"
+ assert variant.runtime_provider_name == "xai_oauth"
+ assert variant.default_model == "grok-4.5"
+ assert variant.models == ("grok-4.5",)
+ assert variant.credential_path
+
+
+@pytest.mark.parametrize("auth_mode", ("api_key", "oauth"))
+@pytest.mark.parametrize(
+ "model_alias",
+ ("grok-4.5", "grok-4.5-latest", "grok-build-latest", "xai/grok-4.5"),
+)
+def test_grok_model_aliases_normalize_to_canonical_id(
+ auth_mode: str, model_alias: str
+) -> None:
+ assert normalize_model_id_for_variant("xai", auth_mode, model_alias) == "grok-4.5"
+
+
+@pytest.mark.parametrize("alias", ("xai", "XAI"))
+def test_grok_runtime_aliases_select_xai(alias: str) -> None:
+ assert normalize_provider_name(alias) == "XAI"
+
+
+@pytest.mark.parametrize("removed_provider", ("grok", "x.ai", "grok_oauth"))
+def test_removed_xai_provider_aliases_are_rejected(removed_provider: str) -> None:
+ with pytest.raises(ValueError, match="Unsupported provider"):
+ load_llm(removed_provider, model="grok-4.5")
+
+
+def test_grok_profile_wires_api_base_context_and_environment_key(monkeypatch) -> None:
+ monkeypatch.setenv("XAI_API_KEY", "xai-env-key")
+ variant = resolve_provider_variant("xai", "api_key")
+ profile = create_profile_for_variant(
+ variant,
+ SetupSelection(
+ family_id="xai",
+ variant_id="XAI",
+ auth_mode="api_key",
+ model="grok-build-latest",
+ api_key_source="env",
+ ),
+ temperature=0.4,
+ )
+
+ assert profile.provider == "XAI"
+ assert profile.provider_family == "xai"
+ assert profile.model == "grok-4.5"
+ assert profile.temperature == 0.4
+ assert profile.base_url == XAI_API_BASE
+ assert profile.api_base == XAI_API_BASE
+ assert profile.kwargs == {"context_window": 500_000}
+
+ load_kwargs = profile.to_load_llm_kwargs()
+ assert load_kwargs["api_key"] == "xai-env-key"
+ assert load_kwargs["api_base"] == XAI_API_BASE
+ assert load_kwargs["context_window"] == 500_000
+
+
+def test_grok_profile_resolves_saved_api_key(monkeypatch, tmp_path) -> None:
+ credential_path = tmp_path / "auth-profiles.json"
+ credential_path.write_text(
+ json.dumps({"apiKeys": {"xai": "xai-saved-key"}}), encoding="utf-8"
+ )
+ monkeypatch.setattr(env_keys, "AUTH_PROFILES_PATH", credential_path)
+ monkeypatch.delenv("XAI_API_KEY", raising=False)
+ variant = resolve_provider_variant("xai", "api_key")
+ profile = create_profile_for_variant(
+ variant,
+ SetupSelection(
+ family_id="xai",
+ variant_id="XAI",
+ auth_mode="api_key",
+ model="grok-4.5",
+ api_key_source="file",
+ ),
+ )
+
+ assert profile.to_load_llm_kwargs()["api_key"] == "xai-saved-key"
+
+
+def test_xai_loader_uses_responses_metadata_and_forces_payload_contract() -> None:
+ llm = load_llm(
+ "XAI",
+ model="grok-4.5-latest",
+ api_key="stub",
+ temperature=0.4,
+ top_p=0.7,
+ store=True,
+ reasoning_options={"effort": "high"},
+ additional_kwargs={
+ "presence_penalty": 0.1,
+ "frequency_penalty": 0.2,
+ "stop": ["done"],
+ },
+ )
+
+ assert type(llm).__name__ == "MobilerunOpenAIResponses"
+ assert llm.model == "grok-4.5"
+ assert llm.api_base == XAI_API_BASE
+ assert llm.metadata.context_window == 500_000
+ assert llm.metadata.is_function_calling_model is True
+ assert llm.reasoning_options is None
+
+ payload = llm._get_model_kwargs(
+ model="caller-selected-model",
+ store=True,
+ temperature=0.3,
+ top_p=0.6,
+ presence_penalty=0.4,
+ frequency_penalty=0.5,
+ stop="stop",
+ reasoning={"effort": "low"},
+ extra_body={
+ "model": "extra-body-model",
+ "store": True,
+ "temperature": 0.2,
+ "top_p": 0.8,
+ "presence_penalty": 0.9,
+ "frequency_penalty": 0.9,
+ "stop": "extra-stop",
+ "reasoning": {"effort": "low"},
+ "metadata": {"safe": "value"},
+ },
+ )
+ assert payload["model"] == "grok-4.5"
+ assert payload["store"] is False
+ assert payload["temperature"] == 0.3
+ assert payload["top_p"] == 0.6
+ assert {
+ "presence_penalty",
+ "frequency_penalty",
+ "stop",
+ "reasoning",
+ }.isdisjoint(payload)
+ assert payload["extra_body"] == {
+ "temperature": 0.2,
+ "top_p": 0.8,
+ "metadata": {"safe": "value"},
+ }
+
+
+def test_xai_loader_pins_catalog_context_metadata() -> None:
+ llm = load_llm(
+ "XAI",
+ model="grok-4.5",
+ api_key="stub",
+ context_window=1,
+ )
+
+ assert llm.metadata.context_window == 500_000
+
+
+def test_xai_sync_chat_pins_final_sdk_wire_body() -> None:
+ requests: list[httpx.Request] = []
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ requests.append(request)
+ return httpx.Response(
+ 200,
+ json=_xai_completed_response(usage=_xai_usage()).model_dump(mode="json"),
+ )
+
+ llm = load_llm(
+ "XAI",
+ model="grok-4.5",
+ api_key="stub",
+ http_client=httpx.Client(transport=httpx.MockTransport(handler)),
+ )
+ try:
+ llm.chat(
+ [ChatMessage(role=MessageRole.USER, content="inspect")],
+ model="caller-selected-model",
+ store=True,
+ reasoning={"effort": "high"},
+ presence_penalty=0.3,
+ extra_body={
+ "model": "extra-body-model",
+ "store": True,
+ "reasoning": {"effort": "low"},
+ "presence_penalty": 0.9,
+ "metadata": {"safe": "value"},
+ },
+ )
+ finally:
+ llm._client.close()
+ asyncio.run(llm._aclient.close())
+
+ assert len(requests) == 1
+ request = requests[0]
+ assert str(request.url) == f"{XAI_API_BASE}/responses"
+ payload = json.loads(request.content)
+ assert payload["model"] == "grok-4.5"
+ assert payload["store"] is False
+ assert payload["metadata"] == {"safe": "value"}
+ assert {"reasoning", "presence_penalty"}.isdisjoint(payload)
+
+
+def test_xai_sync_and_async_chat_send_sanitized_multimodal_tool_payloads() -> None:
+ usage = _xai_usage()
+ response = _xai_completed_response(usage=usage)
+ sync_payload: dict[str, Any] = {}
+ async_payload: dict[str, Any] = {}
+
+ def create_sync(**kwargs: Any) -> Response:
+ sync_payload.update(kwargs)
+ return response
+
+ async def create_async(**kwargs: Any) -> Response:
+ async_payload.update(kwargs)
+ return response
+
+ llm = load_llm("XAI", model="grok-4.5", api_key="stub")
+ llm._client = SimpleNamespace(responses=SimpleNamespace(create=create_sync))
+ llm._aclient = SimpleNamespace(responses=SimpleNamespace(create=create_async))
+ messages = [
+ ChatMessage(
+ role=MessageRole.USER,
+ blocks=[
+ TextBlock(text="inspect"),
+ ImageBlock(image=b"png-bytes", image_mimetype="image/png"),
+ ],
+ )
+ ]
+ call_kwargs = {
+ "temperature": 0.4,
+ "top_p": 0.6,
+ "tools": [
+ {
+ "type": "function",
+ "name": "inspect",
+ "description": "Inspect the image",
+ "parameters": {"type": "object", "properties": {}},
+ }
+ ],
+ "presence_penalty": 0.1,
+ "frequency_penalty": 0.2,
+ "stop": "done",
+ "store": True,
+ "reasoning": {"effort": "high"},
+ }
+
+ sync_result = llm.chat(messages, **dict(call_kwargs))
+ async_result = asyncio.run(llm.achat(messages, **dict(call_kwargs)))
+
+ for result in (sync_result, async_result):
+ tool_calls = llm.get_tool_calls_from_response(result)
+ assert [(call.tool_name, call.tool_kwargs) for call in tool_calls] == [
+ ("inspect", {"value": "ok"})
+ ]
+ usage_result = get_usage_from_response("MobilerunOpenAIResponses", result)
+ assert (
+ usage_result.request_tokens,
+ usage_result.response_tokens,
+ usage_result.total_tokens,
+ ) == (7, 4, 11)
+
+ for payload in (sync_payload, async_payload):
+ assert payload["stream"] is False
+ assert payload["model"] == "grok-4.5"
+ assert payload["temperature"] == 0.4
+ assert payload["top_p"] == 0.6
+ assert payload["store"] is False
+ assert payload["tools"] == call_kwargs["tools"]
+ assert payload["input"][0]["content"][0] == {
+ "type": "input_text",
+ "text": "inspect",
+ }
+ assert payload["input"][0]["content"][1]["type"] == "input_image"
+ assert payload["input"][0]["content"][1]["image_url"].startswith(
+ "data:image/png;base64,"
+ )
+ assert {
+ "presence_penalty",
+ "frequency_penalty",
+ "stop",
+ "reasoning",
+ }.isdisjoint(payload)
+
+
+def test_xai_sync_and_async_stream_preserve_completed_usage() -> None:
+ event = ResponseCompletedEvent(
+ response=_xai_completed_response(usage=_xai_usage()),
+ sequence_number=1,
+ type="response.completed",
+ )
+ sync_payload: dict[str, Any] = {}
+ async_payload: dict[str, Any] = {}
+
+ def create_sync(**kwargs: Any): # type: ignore[no-untyped-def]
+ sync_payload.update(kwargs)
+ return iter((event,))
+
+ async def event_stream(): # type: ignore[no-untyped-def]
+ yield event
+
+ async def create_async(**kwargs: Any): # type: ignore[no-untyped-def]
+ async_payload.update(kwargs)
+ return event_stream()
+
+ llm = load_llm("XAI", model="grok-4.5", api_key="stub")
+ llm._client = SimpleNamespace(responses=SimpleNamespace(create=create_sync))
+ llm._aclient = SimpleNamespace(responses=SimpleNamespace(create=create_async))
+ messages = [ChatMessage(role=MessageRole.USER, content="inspect")]
+ runtime_kwargs = {
+ "temperature": 0.4,
+ "top_p": 0.6,
+ "presence_penalty": 0.1,
+ "store": True,
+ }
+
+ sync_result = list(llm.stream_chat(messages, **dict(runtime_kwargs)))[-1]
+
+ async def collect_async(): # type: ignore[no-untyped-def]
+ return [
+ item
+ async for item in await llm.astream_chat(messages, **dict(runtime_kwargs))
+ ]
+
+ async_result = asyncio.run(collect_async())[-1]
+
+ for result in (sync_result, async_result):
+ usage_result = get_usage_from_response("MobilerunOpenAIResponses", result)
+ assert (
+ usage_result.request_tokens,
+ usage_result.response_tokens,
+ usage_result.total_tokens,
+ ) == (7, 4, 11)
+ assert result.additional_kwargs["usage"].total_tokens == 11
+ for payload in (sync_payload, async_payload):
+ assert payload["stream"] is True
+ assert payload["temperature"] == 0.4
+ assert payload["top_p"] == 0.6
+ assert payload["store"] is False
+ assert "presence_penalty" not in payload
+
+
+def test_xai_structured_predict_sanitizes_sync_and_async_call_kwargs() -> None:
+ from llama_index.core.prompts import PromptTemplate
+ from pydantic import BaseModel
+
+ class StructuredResult(BaseModel):
+ value: str
+
+ sync_payload: dict[str, Any] = {}
+ async_payload: dict[str, Any] = {}
+
+ def parse_sync(**kwargs: Any) -> Any:
+ sync_payload.update(kwargs)
+ return SimpleNamespace(output_parsed=StructuredResult(value="OK"))
+
+ async def parse_async(**kwargs: Any) -> Any:
+ async_payload.update(kwargs)
+ return SimpleNamespace(output_parsed=StructuredResult(value="OK"))
+
+ llm = load_llm("XAI", model="grok-4.5", api_key="stub")
+ llm._client = SimpleNamespace(responses=SimpleNamespace(parse=parse_sync))
+ llm._aclient = SimpleNamespace(responses=SimpleNamespace(parse=parse_async))
+ prompt = PromptTemplate("Return {value}")
+ call_kwargs = {
+ "temperature": 0.4,
+ "top_p": 0.6,
+ "presence_penalty": 0.1,
+ "frequency_penalty": 0.2,
+ "stop": "done",
+ "store": True,
+ "tool_choice": "none",
+ "reasoning": {"effort": "high"},
+ "model": "caller-selected-model",
+ "extra_body": {
+ "model": "extra-body-model",
+ "store": True,
+ "tool_choice": "required",
+ "reasoning": {"effort": "low"},
+ "presence_penalty": 0.9,
+ "metadata": {"safe": "value"},
+ },
+ }
+
+ assert (
+ llm.structured_predict(
+ StructuredResult,
+ prompt,
+ llm_kwargs=dict(call_kwargs),
+ value="OK",
+ ).value
+ == "OK"
+ )
+ assert (
+ asyncio.run(
+ llm.astructured_predict(
+ StructuredResult,
+ prompt,
+ llm_kwargs=dict(call_kwargs),
+ value="OK",
+ )
+ ).value
+ == "OK"
+ )
+
+ for payload in (sync_payload, async_payload):
+ assert payload["temperature"] == 0.4
+ assert payload["top_p"] == 0.6
+ assert payload["store"] is False
+ assert "tool_choice" not in payload
+ assert "model" not in payload["extra_body"]
+ assert payload["extra_body"] == {"metadata": {"safe": "value"}}
+ assert {
+ "presence_penalty",
+ "frequency_penalty",
+ "stop",
+ "reasoning",
+ }.isdisjoint(payload)
+
+
+def test_xai_loader_uses_environment_key_for_direct_runtime(monkeypatch) -> None:
+ monkeypatch.setenv("XAI_API_KEY", "xai-runtime-key")
+
+ llm = load_llm("xai", model="grok-4.5")
+
+ assert llm.api_key == "xai-runtime-key"
+
+
+@pytest.mark.parametrize("alias", ("xai", "XAI"))
+def test_xai_runtime_aliases_default_to_canonical_model(
+ alias: str, monkeypatch
+) -> None:
+ monkeypatch.setenv("XAI_API_KEY", "xai-runtime-key")
+
+ llm = load_llm(alias)
+
+ assert llm.model == "grok-4.5"
+
+
+def test_xai_oauth_runtime_uses_grok_oauth_adapter(tmp_path) -> None:
+ llm = load_llm(
+ "xai_oauth",
+ model="grok-4.5-latest",
+ oauth_access_token="stub",
+ credential_path=str(tmp_path / "auth-profiles.json"),
+ )
+ try:
+ assert type(llm).__name__ == "GrokOAuth"
+ assert llm.model == "grok-4.5"
+ finally:
+ llm._client.close()
+ asyncio.run(llm._aclient.close())
+
+
+def test_xai_loader_does_not_fall_back_to_openai_key(monkeypatch) -> None:
+ monkeypatch.delenv("XAI_API_KEY", raising=False)
+ monkeypatch.setenv("OPENAI_API_KEY", "wrong-provider-key")
+
+ with pytest.raises(ValueError, match="XAI_API_KEY"):
+ load_llm("XAI", model="grok-4.5")
+
+
+@pytest.mark.parametrize(
+ "endpoint_override",
+ (
+ {"api_base": "https://attacker.invalid/v1"},
+ {"base_url": "https://attacker.invalid/v1"},
+ {
+ "api_base": "https://attacker.invalid/v1",
+ "base_url": "https://another-attacker.invalid/v1",
+ },
+ ),
+)
+def test_xai_loader_pins_api_endpoint(endpoint_override: dict[str, str]) -> None:
+ llm = load_llm(
+ "XAI",
+ model="grok-4.5",
+ api_key="xai-secret",
+ **endpoint_override,
+ )
+
+ assert llm.api_base == XAI_API_BASE
diff --git a/tests/test_grok_cli.py b/tests/test_grok_cli.py
new file mode 100644
index 00000000..7e14ddbf
--- /dev/null
+++ b/tests/test_grok_cli.py
@@ -0,0 +1,380 @@
+import json
+from io import StringIO
+from types import SimpleNamespace
+
+import pytest
+from click.testing import CliRunner
+from rich.console import Console
+
+import mobilerun.cli.configure_wizard as configure_wizard
+import mobilerun.cli.main as cli_main
+from mobilerun.agent.utils.oauth.login_timeout import OAuthLoginDeadline
+from mobilerun.cli import oauth_actions
+from mobilerun.cli.configure_wizard import ConfigureWizardCallbacks
+from mobilerun.config_manager import MobileConfig
+
+
+def test_xai_oauth_credentials_are_detected_by_nested_slot(tmp_path) -> None:
+ credential_path = tmp_path / "auth-profiles.json"
+ credential_path.write_text(
+ json.dumps(
+ {
+ "openaiOauth": {"access": "unrelated"},
+ "grokOauth": {"access_token": "grok-access"},
+ }
+ ),
+ encoding="utf-8",
+ )
+
+ assert configure_wizard._oauth_credentials_present(
+ str(credential_path), "xai_oauth"
+ )
+
+ credential_path.write_text(
+ json.dumps({"openaiOauth": {"access": "unrelated"}}),
+ encoding="utf-8",
+ )
+ assert not configure_wizard._oauth_credentials_present(
+ str(credential_path), "xai_oauth"
+ )
+
+
+def test_wizard_prepares_xai_oauth_with_selected_model(tmp_path) -> None:
+ calls: list[dict] = []
+ callbacks = ConfigureWizardCallbacks(
+ run_openai_oauth_login=lambda **kwargs: None,
+ run_anthropic_oauth_login=lambda **kwargs: None,
+ run_gemini_oauth_login=lambda **kwargs: None,
+ run_grok_oauth_login=lambda **kwargs: calls.append(kwargs),
+ )
+
+ configure_wizard._prepare_variant_auth(
+ callbacks=callbacks,
+ variant=SimpleNamespace(id="xai_oauth"),
+ credential_path=str(tmp_path / "auth-profiles.json"),
+ selected_model="grok-4.5",
+ )
+
+ assert calls == [
+ {
+ "credential_path": str(tmp_path / "auth-profiles.json"),
+ "model": "grok-4.5",
+ }
+ ]
+
+
+def test_configure_xai_command_forwards_device_code_options(
+ monkeypatch, tmp_path
+) -> None:
+ calls: list[dict] = []
+ monkeypatch.setattr(
+ cli_main,
+ "_run_grok_oauth_login",
+ lambda **kwargs: calls.append(kwargs),
+ )
+ credential_path = tmp_path / "auth-profiles.json"
+
+ result = CliRunner().invoke(
+ cli_main.cli,
+ [
+ "configure",
+ "xai",
+ "--credential-path",
+ str(credential_path),
+ "--model",
+ "grok-4.5",
+ "--timeout",
+ "12",
+ "--no-browser",
+ "--device-code",
+ ],
+ )
+
+ assert result.exit_code == 0, result.output
+ assert calls == [
+ {
+ "credential_path": str(credential_path),
+ "model": "grok-4.5",
+ "timeout": 12.0,
+ "open_browser": False,
+ "device_code": True,
+ }
+ ]
+
+
+def test_configure_anthropic_forwards_timeout_and_browser_preference(
+ monkeypatch, tmp_path
+) -> None:
+ calls: list[dict] = []
+ monkeypatch.setattr(
+ cli_main,
+ "_run_anthropic_oauth_login",
+ lambda **kwargs: calls.append(kwargs),
+ )
+ credential_path = tmp_path / "auth-profiles.json"
+
+ result = CliRunner().invoke(
+ cli_main.cli,
+ [
+ "configure",
+ "anthropic",
+ "--credential-path",
+ str(credential_path),
+ "--timeout",
+ "12",
+ "--no-browser",
+ ],
+ )
+
+ assert result.exit_code == 0, result.output
+ assert calls == [
+ {
+ "credential_path": str(credential_path),
+ "timeout": 12.0,
+ "open_browser": False,
+ }
+ ]
+
+
+def test_configure_oauth_timeout_must_be_positive_and_finite() -> None:
+ result = CliRunner().invoke(
+ cli_main.cli,
+ ["configure", "xai", "--timeout", "nan"],
+ )
+
+ assert result.exit_code == 2
+ assert "finite number greater than zero" in result.output
+
+
+def test_xai_oauth_action_shares_deadline_and_browser_preference(
+ monkeypatch, tmp_path
+) -> None:
+ observed: dict = {}
+
+ class FakeXAI:
+ def __init__(self, **kwargs): # type: ignore[no-untyped-def]
+ observed["init"] = kwargs
+
+ def login(self, **kwargs): # type: ignore[no-untyped-def]
+ observed["login"] = kwargs
+
+ monkeypatch.setattr(oauth_actions, "GrokOAuth", FakeXAI)
+ oauth_actions.run_grok_oauth_login(
+ str(tmp_path / "auth.json"),
+ "grok-4.5",
+ timeout=12,
+ open_browser=False,
+ device_code=True,
+ )
+
+ deadline = observed["login"]["deadline"]
+ assert isinstance(deadline, OAuthLoginDeadline)
+ assert observed["init"]["timeout"] == 12
+ assert observed["login"]["open_browser"] is False
+ assert observed["login"]["device_code"] is True
+
+
+def test_gemini_empty_entitlement_is_not_persisted(monkeypatch, tmp_path) -> None:
+ persisted: list[OAuthLoginDeadline] = []
+
+ class FakeGemini:
+ def __init__(self, **kwargs):
+ pass
+
+ def login(self, **kwargs):
+ return "access"
+
+ def fetch_available_models(self, **kwargs):
+ return []
+
+ def _persist_credentials(self, *, deadline):
+ persisted.append(deadline)
+
+ monkeypatch.setattr(oauth_actions, "GeminiOAuthCodeAssistLLM", FakeGemini)
+ with pytest.raises(RuntimeError, match="no usable models"):
+ oauth_actions.run_gemini_oauth_login(
+ str(tmp_path / "auth.json"),
+ "gemini-3.5-flash-low",
+ open_browser=False,
+ )
+
+ assert persisted == []
+
+
+def test_configure_help_advertises_only_xai_provider_and_login_command() -> None:
+ runner = CliRunner()
+
+ run_help = runner.invoke(cli_main.cli, ["run", "--help"])
+ configure_help = runner.invoke(cli_main.cli, ["configure", "--help"])
+ xai_help = runner.invoke(cli_main.cli, ["configure", "xai", "--help"])
+ removed_grok_command = runner.invoke(cli_main.cli, ["configure", "grok"])
+
+ assert run_help.exit_code == 0
+ assert "XAI" in run_help.output
+ assert "xai_oauth" not in run_help.output
+ assert "grok_oauth" not in run_help.output
+ assert configure_help.exit_code == 0
+ assert "xai" in configure_help.output.lower()
+ assert "grok" not in configure_help.output.lower()
+ assert xai_help.exit_code == 0
+ assert "--device-code" in xai_help.output
+ assert "xai" in xai_help.output.lower()
+ assert "grok" not in xai_help.output.lower()
+ assert removed_grok_command.exit_code != 0
+
+
+@pytest.mark.parametrize(
+ ("auth_mode", "expected_provider"),
+ (("api_key", "XAI"), ("oauth", "xai_oauth")),
+)
+def test_exact_xai_configure_forms_keep_provider_and_auth_fixed(
+ monkeypatch, auth_mode: str, expected_provider: str
+) -> None:
+ config = MobileConfig()
+ saved_configs: list[MobileConfig] = []
+ login_calls: list[dict] = []
+ model_prompts: list[tuple[tuple[str, ...], str]] = []
+
+ monkeypatch.setattr(configure_wizard.ConfigLoader, "load", lambda: config)
+ monkeypatch.setattr(
+ configure_wizard.ConfigLoader,
+ "save",
+ lambda saved: saved_configs.append(saved),
+ )
+
+ def choose_model(models, *, default_model, allow_back=True): # type: ignore[no-untyped-def]
+ model_prompts.append((tuple(models), default_model))
+ return default_model
+
+ monkeypatch.setattr(configure_wizard, "_prompt_model_choice", choose_model)
+ monkeypatch.setattr(
+ configure_wizard,
+ "_prompt_api_key_for_variant",
+ lambda variant: ("xai-env-key", "env"),
+ )
+ monkeypatch.setattr(
+ configure_wizard,
+ "_oauth_credentials_present",
+ lambda credential_path, variant_id: True,
+ )
+ monkeypatch.setattr(
+ configure_wizard,
+ "_prompt_oauth_credential_action",
+ lambda credential_path: "use_existing",
+ )
+ monkeypatch.setattr(
+ configure_wizard,
+ "select_prompt",
+ lambda *args, **kwargs: pytest.fail(
+ "fixed XAI configure flow unexpectedly reopened the top-level menu"
+ ),
+ )
+ monkeypatch.setattr(
+ cli_main,
+ "_run_grok_oauth_login",
+ lambda **kwargs: login_calls.append(kwargs),
+ )
+
+ result = CliRunner().invoke(
+ cli_main.cli,
+ ["configure", "--provider", "XAI", "--auth-mode", auth_mode],
+ )
+
+ assert result.exit_code == 0, result.output
+ assert saved_configs == [config]
+ assert login_calls == []
+ assert model_prompts == [(("grok-4.5",), "grok-4.5")]
+ assert {
+ (profile.provider, profile.provider_family, profile.auth_mode, profile.model)
+ for profile in config.llm_profiles.values()
+ } == {(expected_provider, "xai", auth_mode, "grok-4.5")}
+ assert "xai_oauth" not in result.output
+ assert "Provider: XAI" in result.output
+
+
+def test_provider_only_flag_does_not_auto_enter_model_flow(monkeypatch) -> None:
+ config = MobileConfig()
+ saved_configs: list[MobileConfig] = []
+ menu_calls: list[str] = []
+
+ monkeypatch.setattr(configure_wizard.ConfigLoader, "load", lambda: config)
+ monkeypatch.setattr(
+ configure_wizard.ConfigLoader,
+ "save",
+ lambda saved: saved_configs.append(saved),
+ )
+ monkeypatch.setattr(
+ configure_wizard,
+ "_prompt_model_choice",
+ lambda *args, **kwargs: pytest.fail(
+ "a provider-only invocation must start at the top-level menu"
+ ),
+ )
+
+ def choose_top_level(message, *args, **kwargs): # type: ignore[no-untyped-def]
+ menu_calls.append(message)
+ return "finish"
+
+ monkeypatch.setattr(configure_wizard, "select_prompt", choose_top_level)
+
+ configure_wizard.run_configure_wizard(
+ Console(file=StringIO(), force_terminal=False),
+ ConfigureWizardCallbacks(
+ run_openai_oauth_login=lambda **kwargs: None,
+ run_anthropic_oauth_login=lambda **kwargs: None,
+ run_gemini_oauth_login=lambda **kwargs: None,
+ ),
+ provider="ollama",
+ auth_mode=None,
+ model=None,
+ api_key=None,
+ base_url=None,
+ )
+
+ assert menu_calls == ["Configure"]
+ assert saved_configs == [config]
+
+
+def test_fixed_provider_and_auth_model_back_returns_to_top_level_once(
+ monkeypatch,
+) -> None:
+ config = MobileConfig()
+ saved_configs: list[MobileConfig] = []
+ model_prompts: list[tuple[tuple[str, ...], str]] = []
+ menu_calls: list[str] = []
+
+ monkeypatch.setattr(configure_wizard.ConfigLoader, "load", lambda: config)
+ monkeypatch.setattr(
+ configure_wizard.ConfigLoader,
+ "save",
+ lambda saved: saved_configs.append(saved),
+ )
+
+ def choose_model(models, *, default_model, allow_back=True): # type: ignore[no-untyped-def]
+ model_prompts.append((tuple(models), default_model))
+ return configure_wizard._BACK
+
+ def choose_top_level(message, *args, **kwargs): # type: ignore[no-untyped-def]
+ menu_calls.append(message)
+ return "finish"
+
+ monkeypatch.setattr(configure_wizard, "_prompt_model_choice", choose_model)
+ monkeypatch.setattr(configure_wizard, "select_prompt", choose_top_level)
+
+ configure_wizard.run_configure_wizard(
+ Console(file=StringIO(), force_terminal=False),
+ ConfigureWizardCallbacks(
+ run_openai_oauth_login=lambda **kwargs: None,
+ run_anthropic_oauth_login=lambda **kwargs: None,
+ run_gemini_oauth_login=lambda **kwargs: None,
+ ),
+ provider="XAI",
+ auth_mode="api_key",
+ model=None,
+ api_key=None,
+ base_url=None,
+ )
+
+ assert model_prompts == [(("grok-4.5",), "grok-4.5")]
+ assert menu_calls == ["Configure"]
+ assert saved_configs == [config]
diff --git a/tests/test_grok_oauth.py b/tests/test_grok_oauth.py
new file mode 100644
index 00000000..b7fb4d7c
--- /dev/null
+++ b/tests/test_grok_oauth.py
@@ -0,0 +1,1645 @@
+from __future__ import annotations
+
+import asyncio
+import io
+import json
+import multiprocessing
+import stat
+import threading
+import time
+from pathlib import Path
+from types import SimpleNamespace
+from urllib.parse import parse_qs, urlparse
+
+import httpx
+import jwt
+import pytest
+from cryptography.hazmat.primitives.asymmetric import ec
+from llama_index.core.base.llms.types import (
+ ChatMessage,
+ ImageBlock,
+ MessageRole,
+ TextBlock,
+ ToolCallBlock,
+)
+
+from mobilerun.agent.usage import get_usage_from_response
+from mobilerun.agent.utils.oauth.grok_oauth_llm import (
+ DEFAULT_GROK_CONTEXT_WINDOW,
+ DEFAULT_GROK_MODEL,
+ DEFAULT_GROK_OAUTH_CLIENT_ID,
+ DEFAULT_GROK_OAUTH_ISSUER,
+ DEFAULT_GROK_OAUTH_PROXY,
+ DEFAULT_GROK_OAUTH_SCOPES,
+ GROK_CLI_COMPAT_VERSION,
+ GROK_CLI_COMPAT_VERSION_HEADER,
+ GrokIDTokenValidator,
+ GrokOAuth,
+ GrokOAuthAuth,
+ GrokOAuthCredentials,
+ GrokOAuthCredentialStore,
+ GrokOAuthError,
+ GrokOAuthReloginRequired,
+ GrokOAuthSessionManager,
+ _parse_callback_query,
+)
+from mobilerun.agent.utils.oauth.login_timeout import OAuthLoginDeadline
+from mobilerun.config_manager import auth_profile_store, env_keys
+from mobilerun.config_manager.auth_profile_store import (
+ AuthProfileFormatError,
+ AuthProfileStore,
+)
+
+
+class _AcceptingIDTokenValidator:
+ def __init__(self) -> None:
+ self.calls: list[tuple[str, str | None]] = []
+
+ def validate(self, token: str, *, nonce: str | None): # type: ignore[no-untyped-def]
+ self.calls.append((token, nonce))
+ return {"sub": "user"}
+
+
+class _FakeClock:
+ def __init__(self) -> None:
+ self.now = 0.0
+ self.sleeps: list[float] = []
+
+ def __call__(self) -> float:
+ return self.now
+
+ def advance(self, seconds: float) -> None:
+ self.now += seconds
+
+ def sleep(self, seconds: float) -> None:
+ self.sleeps.append(seconds)
+ self.advance(seconds)
+
+
+def _request_timeout(request: httpx.Request) -> float:
+ return float(request.extensions["timeout"]["read"])
+
+
+def _credentials(
+ access_token: str = "access-old",
+ refresh_token: str = "refresh-old",
+ *,
+ expires_at_ms: int | None = None,
+) -> GrokOAuthCredentials:
+ return GrokOAuthCredentials(
+ access_token=access_token,
+ refresh_token=refresh_token,
+ expires_at_ms=expires_at_ms or int(time.time() * 1000) + 3_600_000,
+ )
+
+
+def _responses_payload(
+ *,
+ output: list[dict[str, object]] | None = None,
+ input_tokens: int = 8,
+ output_tokens: int = 3,
+) -> dict[str, object]:
+ return {
+ "id": "resp_test",
+ "created_at": int(time.time()),
+ "model": DEFAULT_GROK_MODEL,
+ "object": "response",
+ "output": output or [],
+ "parallel_tool_calls": True,
+ "tool_choice": "auto",
+ "tools": [],
+ "status": "completed",
+ "usage": {
+ "input_tokens": input_tokens,
+ "output_tokens": output_tokens,
+ "total_tokens": input_tokens + output_tokens,
+ "input_tokens_details": {"cached_tokens": 0},
+ "output_tokens_details": {"reasoning_tokens": 0},
+ },
+ }
+
+
+def _responses_sse() -> bytes:
+ completed = _responses_payload(
+ output=[
+ {
+ "id": "msg_test",
+ "type": "message",
+ "role": "assistant",
+ "status": "completed",
+ "content": [
+ {"type": "output_text", "text": "hello", "annotations": []}
+ ],
+ }
+ ],
+ )
+ events = (
+ {
+ "type": "response.output_text.delta",
+ "content_index": 0,
+ "delta": "hel",
+ "item_id": "msg_test",
+ "logprobs": [],
+ "output_index": 0,
+ "sequence_number": 1,
+ },
+ {
+ "type": "response.completed",
+ "response": completed,
+ "sequence_number": 2,
+ },
+ )
+ return (
+ "".join(
+ f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in events
+ )
+ + "data: [DONE]\n\n"
+ ).encode()
+
+
+def _refresh_in_subprocess(path, calls, barrier, results): # type: ignore[no-untyped-def]
+ def handler(_: httpx.Request) -> httpx.Response:
+ with calls.get_lock():
+ calls.value += 1
+ return httpx.Response(
+ 200,
+ json={
+ "access_token": "access-new",
+ "refresh_token": "refresh-rotated",
+ "expires_in": 3600,
+ },
+ )
+
+ try:
+ manager = GrokOAuthSessionManager(
+ credential_store=GrokOAuthCredentialStore(path),
+ http_client=httpx.Client(transport=httpx.MockTransport(handler)),
+ id_token_validator=_AcceptingIDTokenValidator(), # type: ignore[arg-type]
+ )
+ barrier.wait(timeout=10)
+ credentials = manager.get_valid_credentials()
+ results.put(("ok", credentials.access_token))
+ except BaseException as exc:
+ results.put(("error", type(exc).__name__, str(exc)))
+
+
+def test_auth_profile_store_preserves_siblings_and_writes_private_file(tmp_path: Path):
+ path = tmp_path / "auth-profiles.json"
+ path.write_text(json.dumps({"openaiOauth": {"access": "keep"}}))
+
+ AuthProfileStore(path).update_slot("grokOauth", {"accessToken": "secret"})
+
+ payload = json.loads(path.read_text())
+ assert payload["openaiOauth"] == {"access": "keep"}
+ assert payload["grokOauth"] == {"accessToken": "secret"}
+ assert stat.S_IMODE(path.stat().st_mode) == 0o600
+ assert not list(tmp_path.glob(".auth-profiles.json.*.tmp"))
+
+
+def test_auth_profile_store_writes_when_fchmod_is_unavailable(
+ monkeypatch, tmp_path: Path
+):
+ path = tmp_path / "auth-profiles.json"
+ monkeypatch.setattr(auth_profile_store, "_FCHMOD", None)
+
+ AuthProfileStore(path).update_slot("grokOauth", {"accessToken": "secret"})
+
+ assert json.loads(path.read_text()) == {"grokOauth": {"accessToken": "secret"}}
+ assert not list(tmp_path.glob(".auth-profiles.json.*.tmp"))
+
+
+def test_auth_profile_store_rejects_malformed_existing_json(tmp_path: Path):
+ path = tmp_path / "auth-profiles.json"
+ path.write_text("not-json")
+
+ with pytest.raises(AuthProfileFormatError, match="malformed"):
+ AuthProfileStore(path).update_slot("grokOauth", {"accessToken": "new"})
+
+ assert path.read_text() == "not-json"
+
+
+def test_auth_profile_store_replace_failure_preserves_original_and_cleans_temp(
+ monkeypatch, tmp_path: Path
+):
+ path = tmp_path / "auth-profiles.json"
+ original = {"openaiOauth": {"access": "keep"}}
+ path.write_text(json.dumps(original))
+
+ def fail_replace(source: Path, destination: Path) -> None:
+ raise OSError("simulated atomic replace failure")
+
+ monkeypatch.setattr(
+ "mobilerun.config_manager.auth_profile_store.os.replace", fail_replace
+ )
+
+ with pytest.raises(OSError, match="simulated atomic replace failure"):
+ AuthProfileStore(path).update_slot(
+ "grokOauth", {"accessToken": "must-not-be-written"}
+ )
+
+ assert json.loads(path.read_text()) == original
+ assert not list(tmp_path.glob(".auth-profiles.json.*.tmp"))
+
+
+def test_saved_api_keys_use_shared_transaction_and_reject_malformed(
+ monkeypatch, tmp_path: Path
+):
+ path = tmp_path / "auth-profiles.json"
+ monkeypatch.setattr(env_keys, "AUTH_PROFILES_PATH", path)
+ monkeypatch.delenv("XAI_API_KEY", raising=False)
+ path.write_text(json.dumps({"grokOauth": {"accessToken": "keep"}}))
+
+ env_keys.save_env_keys({"xai": "api-key"})
+ payload = json.loads(path.read_text())
+ assert payload["grokOauth"] == {"accessToken": "keep"}
+ assert payload["apiKeys"]["xai"] == "api-key"
+ assert stat.S_IMODE(path.stat().st_mode) == 0o600
+
+ path.write_text("broken")
+ with pytest.raises(AuthProfileFormatError):
+ env_keys.save_env_keys({"xai": "replacement"})
+ assert path.read_text() == "broken"
+
+
+def test_grok_credential_schema_round_trip_does_not_persist_id_token(tmp_path: Path):
+ store = GrokOAuthCredentialStore(tmp_path / "auth-profiles.json")
+ credentials = _credentials()
+ store.save(credentials)
+
+ assert store.load() == credentials
+ raw = json.loads(store.path.read_text())["grokOauth"]
+ assert raw == credentials.to_payload()
+ assert raw["type"] == "oauth"
+ assert raw["provider"] == "xai-grok"
+ assert "idToken" not in raw
+ assert raw["issuer"] == DEFAULT_GROK_OAUTH_ISSUER
+ assert raw["clientId"] == DEFAULT_GROK_OAUTH_CLIENT_ID
+
+
+def test_grok_credentials_reject_unpinned_issuer_and_client():
+ payload = _credentials().to_payload()
+ payload["issuer"] = "https://example.invalid"
+ with pytest.raises(ValueError, match="issuer"):
+ GrokOAuthCredentials.from_payload(payload)
+
+ payload = _credentials().to_payload()
+ payload["provider"] = "other"
+ with pytest.raises(ValueError, match="credential type"):
+ GrokOAuthCredentials.from_payload(payload)
+
+ payload = _credentials().to_payload()
+ payload["clientId"] = "other"
+ with pytest.raises(ValueError, match="clientId"):
+ GrokOAuthCredentials.from_payload(payload)
+
+
+def test_id_token_validator_checks_es256_claims_and_nonce():
+ private_key = ec.generate_private_key(ec.SECP256R1())
+ public_key = private_key.public_key()
+ now = int(time.time())
+ token = jwt.encode(
+ {
+ "sub": "user",
+ "iss": DEFAULT_GROK_OAUTH_ISSUER,
+ "aud": DEFAULT_GROK_OAUTH_CLIENT_ID,
+ "iat": now,
+ "exp": now + 300,
+ "nonce": "expected",
+ },
+ private_key,
+ algorithm="ES256",
+ headers={"kid": "test"},
+ )
+ observed_timeouts: list[float] = []
+ jwks_client = SimpleNamespace(timeout=30.0)
+
+ def signing_key(_token: str): # type: ignore[no-untyped-def]
+ observed_timeouts.append(jwks_client.timeout)
+ return SimpleNamespace(key=public_key)
+
+ jwks_client.get_signing_key_from_jwt = signing_key
+ validator = GrokIDTokenValidator(jwks_client)
+
+ deadline = OAuthLoginDeadline(2)
+ assert (
+ validator.validate(token, nonce="expected", deadline=deadline)["sub"] == "user"
+ )
+ assert observed_timeouts[0] == pytest.approx(2, abs=0.1)
+ assert jwks_client.timeout == 30.0
+ with pytest.raises(jwt.InvalidTokenError, match="nonce"):
+ validator.validate(token, nonce="wrong")
+
+
+def test_id_token_validator_recomputes_timeout_for_each_jwks_fetch(monkeypatch):
+ clock = _FakeClock()
+ client = jwt.PyJWKClient("https://auth.x.ai/.well-known/jwks.json", timeout=30)
+ observed_timeouts: list[float] = []
+
+ def fetch_data(): # type: ignore[no-untyped-def]
+ observed_timeouts.append(client.timeout)
+ clock.advance(1)
+ return {}
+
+ def signing_key(_token: str): # type: ignore[no-untyped-def]
+ client.fetch_data()
+ client.fetch_data()
+ return SimpleNamespace(key="public-key")
+
+ client.fetch_data = fetch_data
+ client.get_signing_key_from_jwt = signing_key
+ monkeypatch.setattr(jwt, "decode", lambda *args, **kwargs: {"nonce": "expected"})
+
+ GrokIDTokenValidator(client).validate(
+ "id-token",
+ nonce="expected",
+ deadline=OAuthLoginDeadline(3, clock=clock, sleeper=clock.sleep),
+ )
+
+ assert observed_timeouts == [3, 2]
+ assert client.timeout == 30
+
+
+def test_legacy_id_token_validator_subclass_keeps_prior_signature(tmp_path: Path):
+ class LegacyValidator(GrokIDTokenValidator):
+ def validate(self, token: str, *, nonce: str | None): # type: ignore[no-untyped-def]
+ return {"sub": "user", "nonce": nonce}
+
+ manager = GrokOAuthSessionManager(
+ credential_store=GrokOAuthCredentialStore(tmp_path / "auth.json"),
+ id_token_validator=LegacyValidator(SimpleNamespace()),
+ )
+
+ credentials = manager._credentials_from_token_response(
+ {"access_token": "access", "id_token": "id-token"},
+ nonce="expected",
+ deadline=OAuthLoginDeadline(1),
+ )
+
+ assert credentials.access_token == "access"
+
+
+def test_id_token_validator_rejects_bad_signature():
+ trusted_key = ec.generate_private_key(ec.SECP256R1())
+ untrusted_key = ec.generate_private_key(ec.SECP256R1())
+ now = int(time.time())
+ token = jwt.encode(
+ {
+ "sub": "user",
+ "iss": DEFAULT_GROK_OAUTH_ISSUER,
+ "aud": DEFAULT_GROK_OAUTH_CLIENT_ID,
+ "iat": now,
+ "exp": now + 300,
+ },
+ untrusted_key,
+ algorithm="ES256",
+ headers={"kid": "test"},
+ )
+ validator = GrokIDTokenValidator(
+ SimpleNamespace(
+ get_signing_key_from_jwt=lambda _: SimpleNamespace(
+ key=trusted_key.public_key()
+ )
+ )
+ )
+
+ with pytest.raises(jwt.InvalidSignatureError):
+ validator.validate(token, nonce=None)
+
+
+@pytest.mark.parametrize(
+ ("claim_overrides", "expected_error"),
+ (
+ ({"iss": "https://issuer.invalid"}, jwt.InvalidIssuerError),
+ ({"aud": "different-client"}, jwt.InvalidAudienceError),
+ ({"iat": 1, "exp": 2}, jwt.ExpiredSignatureError),
+ ),
+)
+def test_id_token_validator_rejects_invalid_registered_claims(
+ claim_overrides: dict[str, object],
+ expected_error: type[jwt.InvalidTokenError],
+):
+ private_key = ec.generate_private_key(ec.SECP256R1())
+ now = int(time.time())
+ claims: dict[str, object] = {
+ "sub": "user",
+ "iss": DEFAULT_GROK_OAUTH_ISSUER,
+ "aud": DEFAULT_GROK_OAUTH_CLIENT_ID,
+ "iat": now,
+ "exp": now + 300,
+ }
+ claims.update(claim_overrides)
+ token = jwt.encode(
+ claims,
+ private_key,
+ algorithm="ES256",
+ headers={"kid": "test"},
+ )
+ validator = GrokIDTokenValidator(
+ SimpleNamespace(
+ get_signing_key_from_jwt=lambda _: SimpleNamespace(
+ key=private_key.public_key()
+ )
+ )
+ )
+
+ with pytest.raises(expected_error):
+ validator.validate(token, nonce=None)
+
+
+def test_authorization_url_is_pinned_pkce_and_complete():
+ url = GrokOAuth._build_auth_url(
+ redirect_uri="http://127.0.0.1:54321/callback",
+ code_challenge="challenge",
+ state="state",
+ nonce="nonce",
+ )
+ parsed = urlparse(url)
+ query = parse_qs(parsed.query)
+
+ assert f"{parsed.scheme}://{parsed.netloc}" == DEFAULT_GROK_OAUTH_ISSUER
+ assert parsed.path == "/oauth2/authorize"
+ assert query["client_id"] == [DEFAULT_GROK_OAUTH_CLIENT_ID]
+ assert query["redirect_uri"] == ["http://127.0.0.1:54321/callback"]
+ assert query["scope"] == [" ".join(DEFAULT_GROK_OAUTH_SCOPES)]
+ assert query["code_challenge_method"] == ["S256"]
+ assert query["code_challenge"] == ["challenge"]
+ assert query["state"] == ["state"]
+ assert query["nonce"] == ["nonce"]
+ assert query["referrer"] == ["grok-build"]
+
+
+def test_callback_query_rejects_duplicate_code_state_and_error_values():
+ assert _parse_callback_query("code=one&state=expected") == {
+ "code": "one",
+ "state": "expected",
+ "error": None,
+ }
+ for query in (
+ "code=one&code=two&state=expected",
+ "code=&code=one&state=expected",
+ "code=one&state=a&state=b",
+ "code=one&state=&state=expected",
+ "error=denied&error=other",
+ ):
+ parsed = _parse_callback_query(query)
+ assert parsed["error"] == "invalid_callback"
+
+
+def test_browser_login_uses_random_loopback_callback_and_exchanges_code(
+ monkeypatch, tmp_path: Path
+):
+ clock = _FakeClock()
+ deadline = OAuthLoginDeadline(5, clock=clock, sleeper=clock.sleep)
+ token_requests: list[httpx.Request] = []
+
+ def token_handler(request: httpx.Request) -> httpx.Response:
+ token_requests.append(request)
+ return httpx.Response(
+ 200,
+ json={
+ "access_token": "access",
+ "refresh_token": "refresh",
+ "expires_in": 3600,
+ "id_token": "signed-id-token",
+ },
+ )
+
+ manager = GrokOAuthSessionManager(
+ credential_store=GrokOAuthCredentialStore(tmp_path / "auth.json"),
+ http_client=httpx.Client(transport=httpx.MockTransport(token_handler)),
+ id_token_validator=_AcceptingIDTokenValidator(), # type: ignore[arg-type]
+ )
+ llm = GrokOAuth(oauth_session_manager=manager)
+ opened: list[str] = []
+ callback_handler: list[type] = []
+
+ class FakeRequest:
+ def __init__(self, target: str) -> None:
+ self.input = io.BytesIO(
+ f"GET {target} HTTP/1.1\r\nHost: 127.0.0.1\r\n\r\n".encode()
+ )
+ self.output = io.BytesIO()
+
+ def makefile(self, mode: str, *args, **kwargs): # type: ignore[no-untyped-def]
+ return self.input if "r" in mode else self.output
+
+ def sendall(self, data: bytes) -> None:
+ self.output.write(data)
+
+ def close(self) -> None:
+ return
+
+ class FakeHTTPServer:
+ server_name = "127.0.0.1"
+ server_port = 54321
+
+ def __init__(self, address, handler): # type: ignore[no-untyped-def]
+ assert address == ("127.0.0.1", 0)
+ self.server_address = ("127.0.0.1", self.server_port)
+ callback_handler.append(handler)
+
+ def serve_forever(self) -> None:
+ return
+
+ def shutdown(self) -> None:
+ return
+
+ def server_close(self) -> None:
+ return
+
+ def complete_in_browser(authorization_url: str) -> bool:
+ opened.append(authorization_url)
+ clock.advance(1)
+ query = parse_qs(urlparse(authorization_url).query)
+ redirect_uri = query["redirect_uri"][0]
+ assert redirect_uri.startswith("http://127.0.0.1:")
+ assert redirect_uri.endswith("/callback")
+ assert ":0/" not in redirect_uri
+ callback_url = f"{redirect_uri}?code=code&state={query['state'][0]}"
+ callback_target = urlparse(callback_url)
+ request = FakeRequest(f"{callback_target.path}?{callback_target.query}")
+ callback_handler[0](request, ("127.0.0.1", 12345), fake_server[0])
+ assert b"200" in request.output.getvalue().splitlines()[0]
+ return True
+
+ fake_server: list[FakeHTTPServer] = []
+
+ def make_server(address, handler): # type: ignore[no-untyped-def]
+ server = FakeHTTPServer(address, handler)
+ fake_server.append(server)
+ return server
+
+ monkeypatch.setattr(
+ "mobilerun.agent.utils.oauth.grok_oauth_llm.HTTPServer", make_server
+ )
+ monkeypatch.setattr("webbrowser.open", complete_in_browser)
+ credentials = llm.login(
+ open_browser=True,
+ timeout_seconds=5,
+ deadline=deadline,
+ )
+
+ assert credentials.access_token == "access"
+ assert len(opened) == 1
+ token_form = parse_qs(token_requests[0].content.decode())
+ assert token_form["code"] == ["code"]
+ assert token_form["redirect_uri"][0].startswith("http://127.0.0.1:")
+ assert token_form["code_verifier"][0]
+ assert _request_timeout(token_requests[0]) == pytest.approx(4)
+
+
+def test_code_exchange_uses_form_validates_id_and_saves_only_session(tmp_path: Path):
+ requests: list[httpx.Request] = []
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ requests.append(request)
+ return httpx.Response(
+ 200,
+ json={
+ "access_token": "access-new",
+ "refresh_token": "refresh-new",
+ "expires_in": 3600,
+ "token_type": "Bearer",
+ "scope": "openid offline_access api:access",
+ "id_token": "signed-id-token",
+ },
+ )
+
+ validator = _AcceptingIDTokenValidator()
+ manager = GrokOAuthSessionManager(
+ credential_store=GrokOAuthCredentialStore(tmp_path / "auth.json"),
+ http_client=httpx.Client(transport=httpx.MockTransport(handler)),
+ id_token_validator=validator, # type: ignore[arg-type]
+ )
+ credentials = manager.exchange_authorization_code(
+ code="auth-code",
+ redirect_uri="http://127.0.0.1:54321/callback",
+ code_verifier="verifier",
+ nonce="nonce",
+ )
+
+ form = parse_qs(requests[0].content.decode())
+ assert form["grant_type"] == ["authorization_code"]
+ assert form["client_id"] == [DEFAULT_GROK_OAUTH_CLIENT_ID]
+ assert form["code_verifier"] == ["verifier"]
+ assert requests[0].headers["Accept"] == "application/json"
+ assert validator.calls == [("signed-id-token", "nonce")]
+ assert credentials.refresh_token == "refresh-new"
+ raw = json.loads((tmp_path / "auth.json").read_text())["grokOauth"]
+ assert "signed-id-token" not in json.dumps(raw)
+
+
+def test_non_browser_token_accepts_missing_id_token_and_uses_access_jwt_expiry(
+ tmp_path: Path,
+):
+ now = int(time.time())
+ # The access-token claim is used only as refresh scheduling metadata.
+ unsigned_access = jwt.encode(
+ {"exp": now + 900}, "not-a-provider-key-that-is-long-enough", algorithm="HS256"
+ )
+
+ validator = _AcceptingIDTokenValidator()
+ manager = GrokOAuthSessionManager(
+ credential_store=GrokOAuthCredentialStore(tmp_path / "auth.json"),
+ id_token_validator=validator, # type: ignore[arg-type]
+ )
+
+ credentials = manager._credentials_from_token_response(
+ {
+ "access_token": unsigned_access,
+ "refresh_token": "refresh",
+ },
+ nonce=None,
+ )
+
+ assert credentials.expires_at_ms == (now + 900) * 1000
+ assert validator.calls == []
+
+
+def test_browser_code_exchange_rejects_missing_id_token(tmp_path: Path):
+ manager = GrokOAuthSessionManager(
+ credential_store=GrokOAuthCredentialStore(tmp_path / "auth.json"),
+ http_client=httpx.Client(
+ transport=httpx.MockTransport(
+ lambda _: httpx.Response(
+ 200,
+ json={
+ "access_token": "access",
+ "refresh_token": "refresh",
+ "expires_in": 3600,
+ },
+ )
+ )
+ ),
+ id_token_validator=_AcceptingIDTokenValidator(), # type: ignore[arg-type]
+ )
+
+ with pytest.raises(GrokOAuthError, match="did not contain an ID token"):
+ manager.exchange_authorization_code(
+ code="code",
+ redirect_uri="http://127.0.0.1:54321/callback",
+ code_verifier="verifier",
+ nonce="nonce",
+ )
+
+ assert manager.credential_store.load() is None
+
+
+@pytest.mark.parametrize(
+ ("status_code", "body"),
+ (
+ (200, "secret-token-not-json"),
+ (400, '{"error_description":"secret-token"}'),
+ ),
+)
+def test_token_errors_never_expose_response_bodies(
+ status_code: int, body: str, tmp_path: Path
+):
+ manager = GrokOAuthSessionManager(
+ credential_store=GrokOAuthCredentialStore(tmp_path / "auth.json"),
+ http_client=httpx.Client(
+ transport=httpx.MockTransport(
+ lambda _: httpx.Response(status_code, text=body)
+ )
+ ),
+ id_token_validator=_AcceptingIDTokenValidator(), # type: ignore[arg-type]
+ )
+
+ with pytest.raises(GrokOAuthError) as captured:
+ manager.exchange_authorization_code(
+ code="code",
+ redirect_uri="http://127.0.0.1:54321/callback",
+ code_verifier="verifier",
+ nonce="nonce",
+ )
+
+ assert "secret-token" not in str(captured.value)
+
+
+def test_authorization_code_timeout_is_not_retried(tmp_path: Path):
+ attempts = 0
+ delays: list[float] = []
+ request_timeouts: list[float] = []
+ deadline = OAuthLoginDeadline(2)
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ nonlocal attempts
+ attempts += 1
+ request_timeouts.append(_request_timeout(request))
+ raise httpx.ReadTimeout("ambiguous token exchange timeout", request=request)
+
+ manager = GrokOAuthSessionManager(
+ credential_store=GrokOAuthCredentialStore(tmp_path / "auth.json"),
+ http_client=httpx.Client(transport=httpx.MockTransport(handler)),
+ id_token_validator=_AcceptingIDTokenValidator(), # type: ignore[arg-type]
+ request_timeout=20,
+ retry_backoff_seconds=(0.01, 0.02),
+ sleep=delays.append,
+ )
+
+ with pytest.raises(GrokOAuthError, match="transient network error") as captured:
+ manager.exchange_authorization_code(
+ code="one-time-code",
+ redirect_uri="http://127.0.0.1:54321/callback",
+ code_verifier="verifier",
+ nonce="nonce",
+ deadline=deadline,
+ )
+
+ assert attempts == 1
+ assert request_timeouts == pytest.approx([2], abs=0.1)
+ assert delays == []
+ assert "one-time-code" not in str(captured.value)
+ assert manager.credential_store.load() is None
+
+
+def test_device_flow_matches_xai_surface_and_handles_pending(
+ monkeypatch, capsys, tmp_path: Path
+):
+ requests: list[httpx.Request] = []
+ verification_uri = "https://accounts.x.ai/device?code=ABCD-1234"
+ opened: list[str] = []
+ opened_event = threading.Event()
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ requests.append(request)
+ if request.url.path.endswith("/device/code"):
+ return httpx.Response(
+ 200,
+ json={
+ "device_code": "device-secret",
+ "user_code": "ABCD-1234",
+ "verification_uri_complete": verification_uri,
+ "expires_in": 1800,
+ "interval": 1,
+ },
+ )
+ if len([r for r in requests if r.url.path.endswith("/token")]) == 1:
+ return httpx.Response(400, json={"error": "authorization_pending"})
+ return httpx.Response(
+ 200,
+ json={
+ "access_token": "access",
+ "refresh_token": "refresh",
+ "expires_in": 3600,
+ },
+ )
+
+ monkeypatch.setattr(time, "sleep", lambda _: None)
+ monkeypatch.setattr(
+ "webbrowser.open",
+ lambda url: (opened.append(url), opened_event.set()),
+ )
+ manager = GrokOAuthSessionManager(
+ credential_store=GrokOAuthCredentialStore(tmp_path / "auth.json"),
+ http_client=httpx.Client(transport=httpx.MockTransport(handler)),
+ id_token_validator=_AcceptingIDTokenValidator(), # type: ignore[arg-type]
+ )
+ llm = GrokOAuth(oauth_session_manager=manager)
+
+ credentials = llm.login(
+ device_code=True,
+ open_browser=True,
+ timeout_seconds=10,
+ )
+
+ device_request = requests[0]
+ assert device_request.headers["x-grok-client-surface"] == "grok-build"
+ assert device_request.headers["Accept"] == "application/json"
+ assert parse_qs(device_request.content.decode()) == {
+ "client_id": [DEFAULT_GROK_OAUTH_CLIENT_ID],
+ "scope": [" ".join(DEFAULT_GROK_OAUTH_SCOPES)],
+ }
+ token_form = parse_qs(requests[1].content.decode())
+ assert token_form["device_code"] == ["device-secret"]
+ assert "user_code" not in token_form
+ assert credentials.access_token == "access"
+ assert manager.credential_store.load() == credentials
+ assert opened_event.wait(timeout=1)
+ assert opened == [verification_uri]
+ assert verification_uri in capsys.readouterr().out
+
+
+def test_bind_fallback_preserves_deadline_and_no_browser(monkeypatch, tmp_path: Path):
+ manager = GrokOAuthSessionManager(
+ credential_store=GrokOAuthCredentialStore(tmp_path / "auth.json")
+ )
+ llm = GrokOAuth(oauth_session_manager=manager)
+ deadline = OAuthLoginDeadline(2)
+ expected = _credentials()
+ calls: list[dict[str, object]] = []
+
+ monkeypatch.setattr(
+ "mobilerun.agent.utils.oauth.grok_oauth_llm._is_headless_environment",
+ lambda: False,
+ )
+ monkeypatch.setattr(
+ "mobilerun.agent.utils.oauth.grok_oauth_llm.HTTPServer",
+ lambda *args, **kwargs: (_ for _ in ()).throw(OSError("unavailable")),
+ )
+ monkeypatch.setattr(
+ llm,
+ "_login_device_code",
+ lambda **kwargs: (calls.append(kwargs), expected)[1],
+ )
+
+ assert llm.login(open_browser=False, deadline=deadline) == expected
+ assert calls == [{"deadline": deadline, "open_browser": False}]
+
+
+def test_device_poll_retries_shrink_deadline_and_do_not_save_after_expiry(
+ tmp_path: Path,
+):
+ clock = _FakeClock()
+ deadline = OAuthLoginDeadline(2, clock=clock, sleeper=clock.sleep)
+ poll_attempts = 0
+ request_timeouts: list[float] = []
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ nonlocal poll_attempts
+ request_timeouts.append(_request_timeout(request))
+ if request.url.path.endswith("/device/code"):
+ clock.advance(0.4)
+ return httpx.Response(
+ 200,
+ json={
+ "device_code": "device-secret",
+ "user_code": "ABCD-1234",
+ "verification_uri": "https://accounts.x.ai/device",
+ "expires_in": 1800,
+ "interval": 1,
+ },
+ )
+ poll_attempts += 1
+ if poll_attempts == 1:
+ clock.advance(0.3)
+ raise httpx.ConnectError("temporary connect failure", request=request)
+ if poll_attempts == 2:
+ clock.advance(0.2)
+ return httpx.Response(503, text="sensitive-upstream-body")
+ clock.advance(0.36)
+ return httpx.Response(
+ 200,
+ json={
+ "access_token": "access",
+ "refresh_token": "refresh",
+ "expires_in": 3600,
+ },
+ )
+
+ store = GrokOAuthCredentialStore(tmp_path / "auth.json")
+ manager = GrokOAuthSessionManager(
+ credential_store=store,
+ http_client=httpx.Client(transport=httpx.MockTransport(handler)),
+ id_token_validator=_AcceptingIDTokenValidator(), # type: ignore[arg-type]
+ request_timeout=20,
+ retry_backoff_seconds=(0.25, 0.5),
+ )
+
+ with pytest.raises(TimeoutError, match="OAuth login timed out"):
+ GrokOAuth(oauth_session_manager=manager).login(
+ device_code=True,
+ open_browser=False,
+ deadline=deadline,
+ )
+
+ assert poll_attempts == 3
+ assert request_timeouts == pytest.approx([2.0, 1.6, 1.05, 0.35])
+ assert clock.sleeps == pytest.approx([0.25, 0.5])
+ assert store.load() is None
+
+
+def test_refresh_preserves_rotated_token_and_is_coordinated_across_managers(
+ tmp_path: Path,
+):
+ path = tmp_path / "auth.json"
+ store = GrokOAuthCredentialStore(path)
+ store.save(_credentials(expires_at_ms=1))
+ calls = 0
+ calls_lock = threading.Lock()
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ nonlocal calls
+ with calls_lock:
+ calls += 1
+ return httpx.Response(
+ 200,
+ json={
+ "access_token": "access-new",
+ "refresh_token": "refresh-rotated",
+ "expires_in": 3600,
+ },
+ )
+
+ managers = [
+ GrokOAuthSessionManager(
+ credential_store=GrokOAuthCredentialStore(path),
+ http_client=httpx.Client(transport=httpx.MockTransport(handler)),
+ id_token_validator=_AcceptingIDTokenValidator(), # type: ignore[arg-type]
+ )
+ for _ in range(2)
+ ]
+ results: list[GrokOAuthCredentials] = []
+ threads = [
+ threading.Thread(
+ target=lambda manager=m: results.append(manager.get_valid_credentials())
+ )
+ for m in managers
+ ]
+ for thread in threads:
+ thread.start()
+ for thread in threads:
+ thread.join()
+
+ assert calls == 1
+ assert len(results) == 2
+ assert {result.access_token for result in results} == {"access-new"}
+ assert store.load().refresh_token == "refresh-rotated" # type: ignore[union-attr]
+
+
+def test_refresh_retries_timeout_and_5xx_then_persists_rotation(tmp_path: Path):
+ path = tmp_path / "auth.json"
+ store = GrokOAuthCredentialStore(path)
+ store.save(_credentials(expires_at_ms=1))
+ attempts = 0
+ delays: list[float] = []
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ nonlocal attempts
+ attempts += 1
+ if attempts == 1:
+ raise httpx.ReadTimeout("temporary timeout", request=request)
+ if attempts == 2:
+ return httpx.Response(502, text="sensitive-upstream-body")
+ return httpx.Response(
+ 200,
+ json={
+ "access_token": "access-new",
+ "refresh_token": "refresh-rotated",
+ "expires_in": 3600,
+ },
+ )
+
+ manager = GrokOAuthSessionManager(
+ credential_store=store,
+ http_client=httpx.Client(transport=httpx.MockTransport(handler)),
+ id_token_validator=_AcceptingIDTokenValidator(), # type: ignore[arg-type]
+ retry_backoff_seconds=(0.01, 0.02),
+ sleep=delays.append,
+ )
+
+ credentials = manager.get_valid_credentials()
+
+ assert attempts == 3
+ assert delays == [0.01, 0.02]
+ assert credentials.access_token == "access-new"
+ assert store.load().refresh_token == "refresh-rotated" # type: ignore[union-attr]
+
+
+@pytest.mark.parametrize("oauth_error", ("invalid_grant", "invalid_client"))
+def test_permanent_refresh_rejection_requires_relogin_without_changing_disk(
+ oauth_error: str, tmp_path: Path
+):
+ path = tmp_path / "auth.json"
+ store = GrokOAuthCredentialStore(path)
+ store.save(_credentials(expires_at_ms=1))
+ store.profile_store.update_slot("openaiOauth", {"access": "keep-sibling"})
+ before = path.read_bytes()
+
+ manager = GrokOAuthSessionManager(
+ credential_store=store,
+ http_client=httpx.Client(
+ transport=httpx.MockTransport(
+ lambda _: httpx.Response(
+ 400,
+ json={
+ "error": oauth_error,
+ "error_description": "sensitive-refresh-response-body",
+ },
+ )
+ )
+ ),
+ id_token_validator=_AcceptingIDTokenValidator(), # type: ignore[arg-type]
+ sleep=lambda _: None,
+ )
+ manager._credentials = _credentials(expires_at_ms=1)
+
+ with pytest.raises(
+ GrokOAuthReloginRequired, match="re-login is required"
+ ) as captured:
+ manager.get_valid_credentials()
+
+ assert "sensitive-refresh-response-body" not in str(captured.value)
+ assert manager._credentials is None
+ assert path.read_bytes() == before
+ profile = json.loads(path.read_text())
+ assert profile["openaiOauth"] == {"access": "keep-sibling"}
+ assert profile["grokOauth"]["accessToken"] == "access-old"
+
+
+def test_refresh_is_coordinated_across_processes(tmp_path: Path):
+ path = tmp_path / "auth.json"
+ store = GrokOAuthCredentialStore(path)
+ store.save(_credentials(expires_at_ms=1))
+ context = multiprocessing.get_context("spawn")
+ calls = context.Value("i", 0)
+ barrier = context.Barrier(2)
+ results = context.Queue()
+ processes = [
+ context.Process(
+ target=_refresh_in_subprocess,
+ args=(str(path), calls, barrier, results),
+ )
+ for _ in range(2)
+ ]
+
+ try:
+ for process in processes:
+ process.start()
+ received = [results.get(timeout=15) for _ in processes]
+ for process in processes:
+ process.join(timeout=15)
+ finally:
+ for process in processes:
+ if process.is_alive():
+ process.terminate()
+ process.join(timeout=5)
+
+ assert received == [("ok", "access-new"), ("ok", "access-new")]
+ assert [process.exitcode for process in processes] == [0, 0]
+ assert calls.value == 1
+ assert store.load().refresh_token == "refresh-rotated" # type: ignore[union-attr]
+
+
+class _StubSessionManager:
+ def __init__(self) -> None:
+ self.calls: list[dict[str, object]] = []
+
+ def get_valid_credentials(self, **kwargs): # type: ignore[no-untyped-def]
+ self.calls.append(kwargs)
+ if kwargs.get("force_refresh"):
+ return _credentials("access-new", "refresh-new")
+ return _credentials()
+
+
+def test_sync_auth_injects_headers_and_retries_exactly_one_401():
+ manager = _StubSessionManager()
+ seen: list[str] = []
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ seen.append(request.headers["Authorization"])
+ assert request.headers["X-XAI-Token-Auth"] == "xai-grok-cli"
+ assert request.headers["x-grok-model-override"] == DEFAULT_GROK_MODEL
+ assert (
+ request.headers[GROK_CLI_COMPAT_VERSION_HEADER] == GROK_CLI_COMPAT_VERSION
+ )
+ return httpx.Response(401 if len(seen) == 1 else 200, json={})
+
+ with httpx.Client(
+ transport=httpx.MockTransport(handler),
+ auth=GrokOAuthAuth(manager, model=DEFAULT_GROK_MODEL), # type: ignore[arg-type]
+ ) as client:
+ response = client.post("https://example.test/responses", json={"input": "hi"})
+
+ assert response.status_code == 200
+ assert seen == ["Bearer access-old", "Bearer access-new"]
+ assert manager.calls == [
+ {},
+ {"force_refresh": True, "rejected_access_token": "access-old"},
+ ]
+
+
+def test_sync_auth_replays_only_once_on_consecutive_401s():
+ manager = _StubSessionManager()
+ seen: list[str] = []
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ seen.append(request.headers["Authorization"])
+ return httpx.Response(401, json={})
+
+ with httpx.Client(
+ transport=httpx.MockTransport(handler),
+ auth=GrokOAuthAuth(manager, model=DEFAULT_GROK_MODEL), # type: ignore[arg-type]
+ ) as client:
+ response = client.post("https://example.test/responses", json={"input": "hi"})
+
+ assert response.status_code == 401
+ assert seen == ["Bearer access-old", "Bearer access-new"]
+ assert manager.calls == [
+ {},
+ {"force_refresh": True, "rejected_access_token": "access-old"},
+ ]
+
+
+@pytest.mark.parametrize("status_code", (400, 403, 429))
+def test_sync_auth_does_not_replay_non_401_responses(status_code: int):
+ manager = _StubSessionManager()
+ seen: list[str] = []
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ seen.append(request.headers["Authorization"])
+ return httpx.Response(status_code, json={})
+
+ with httpx.Client(
+ transport=httpx.MockTransport(handler),
+ auth=GrokOAuthAuth(manager, model=DEFAULT_GROK_MODEL), # type: ignore[arg-type]
+ ) as client:
+ response = client.post("https://example.test/responses", json={"input": "hi"})
+
+ assert response.status_code == status_code
+ assert seen == ["Bearer access-old"]
+ assert manager.calls == [{}]
+
+
+def test_streaming_request_hides_initial_401_and_replays_before_body_exposure():
+ manager = _StubSessionManager()
+ seen: list[str] = []
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ seen.append(request.headers["Authorization"])
+ if len(seen) == 1:
+ return httpx.Response(401, content=b"rejected-stream-body")
+ return httpx.Response(200, content=b"data: response-event\n\n")
+
+ with httpx.Client(
+ transport=httpx.MockTransport(handler),
+ auth=GrokOAuthAuth(manager, model=DEFAULT_GROK_MODEL), # type: ignore[arg-type]
+ ) as client:
+ with client.stream(
+ "POST", "https://example.test/responses", json={"stream": True}
+ ) as response:
+ exposed_body = b"".join(response.iter_bytes())
+
+ assert response.status_code == 200
+ assert exposed_body == b"data: response-event\n\n"
+ assert b"rejected" not in exposed_body
+ assert seen == ["Bearer access-old", "Bearer access-new"]
+
+
+def test_async_auth_injects_headers_and_retries_exactly_one_401():
+ manager = _StubSessionManager()
+ seen: list[str] = []
+
+ async def run() -> httpx.Response:
+ def handler(request: httpx.Request) -> httpx.Response:
+ seen.append(request.headers["Authorization"])
+ assert request.headers["X-XAI-Token-Auth"] == "xai-grok-cli"
+ assert request.headers["x-grok-model-override"] == DEFAULT_GROK_MODEL
+ assert (
+ request.headers[GROK_CLI_COMPAT_VERSION_HEADER]
+ == GROK_CLI_COMPAT_VERSION
+ )
+ return httpx.Response(401 if len(seen) == 1 else 200, json={})
+
+ async with httpx.AsyncClient(
+ transport=httpx.MockTransport(handler),
+ auth=GrokOAuthAuth(manager, model=DEFAULT_GROK_MODEL), # type: ignore[arg-type]
+ ) as client:
+ return await client.post(
+ "https://example.test/responses", json={"input": "hi"}
+ )
+
+ response = asyncio.run(run())
+ assert response.status_code == 200
+ assert seen == ["Bearer access-old", "Bearer access-new"]
+ assert manager.calls == [
+ {},
+ {"force_refresh": True, "rejected_access_token": "access-old"},
+ ]
+
+
+def test_oauth_adapter_sync_chat_serializes_image_and_tool_on_sanitized_wire(
+ tmp_path: Path,
+):
+ requests: list[httpx.Request] = []
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ requests.append(request)
+ return httpx.Response(
+ 200,
+ json=_responses_payload(
+ output=[
+ {
+ "id": "msg_test",
+ "type": "message",
+ "role": "assistant",
+ "status": "completed",
+ "content": [
+ {
+ "type": "output_text",
+ "text": "I will tap.",
+ "annotations": [],
+ }
+ ],
+ },
+ {
+ "id": "call_item",
+ "call_id": "call_test",
+ "type": "function_call",
+ "name": "tap",
+ "arguments": '{"x":1}',
+ "status": "completed",
+ },
+ ]
+ ),
+ )
+
+ llm = GrokOAuth(
+ oauth_credential_path=str(tmp_path / "auth.json"),
+ oauth_access_token="adapter-access",
+ oauth_refresh_token="adapter-refresh",
+ oauth_expires_at_ms=int(time.time() * 1000) + 3_600_000,
+ http_client=httpx.Client(transport=httpx.MockTransport(handler)),
+ )
+ try:
+ response = llm.chat(
+ [
+ ChatMessage(
+ role=MessageRole.USER,
+ blocks=[
+ TextBlock(text="inspect this image"),
+ ImageBlock(image=b"\x89PNG\r\n\x1a\n"),
+ ],
+ )
+ ],
+ tools=[
+ {
+ "type": "function",
+ "name": "tap",
+ "description": "Tap a coordinate",
+ "parameters": {"type": "object"},
+ }
+ ],
+ tool_choice={"type": "function", "name": "tap"},
+ temperature=0.2,
+ top_p=0.7,
+ presence_penalty=0.3,
+ frequency_penalty=0.4,
+ stop="done",
+ reasoning={"effort": "high"},
+ store=True,
+ model="caller-selected-model",
+ extra_body={
+ "model": "extra-body-model",
+ "store": True,
+ "temperature": 0.9,
+ "top_p": 0.8,
+ "presence_penalty": 0.7,
+ "frequency_penalty": 0.6,
+ "stop": "extra-stop",
+ "reasoning": {"effort": "low"},
+ "metadata": {"safe": "value"},
+ },
+ )
+ finally:
+ llm._client.close()
+ asyncio.run(llm._aclient.close())
+
+ assert response.message.content == "I will tap."
+ tool_call = next(
+ block for block in response.message.blocks if isinstance(block, ToolCallBlock)
+ )
+ assert tool_call.tool_name == "tap"
+ assert json.loads(tool_call.tool_kwargs) == {"x": 1}
+ assert response.additional_kwargs["usage"].total_tokens == 11
+
+ assert len(requests) == 1
+ request = requests[0]
+ assert str(request.url) == f"{DEFAULT_GROK_OAUTH_PROXY}/responses"
+ assert request.headers["Authorization"] == "Bearer adapter-access"
+ assert request.headers["X-XAI-Token-Auth"] == "xai-grok-cli"
+ assert request.headers["x-grok-model-override"] == DEFAULT_GROK_MODEL
+ assert request.headers[GROK_CLI_COMPAT_VERSION_HEADER] == GROK_CLI_COMPAT_VERSION
+ payload = json.loads(request.content)
+ assert payload["model"] == DEFAULT_GROK_MODEL
+ assert payload["store"] is False
+ assert payload["metadata"] == {"safe": "value"}
+ assert payload["stream"] is False
+ assert payload["tool_choice"] == {"type": "function", "name": "tap"}
+ assert payload["tools"][0]["name"] == "tap"
+ content = payload["input"][0]["content"]
+ assert content[0] == {"type": "input_text", "text": "inspect this image"}
+ assert content[1]["type"] == "input_image"
+ assert content[1]["image_url"].startswith("data:image/png;base64,")
+ for unsupported in (
+ "temperature",
+ "top_p",
+ "presence_penalty",
+ "frequency_penalty",
+ "stop",
+ "reasoning",
+ ):
+ assert unsupported not in payload
+
+
+def test_oauth_adapter_sync_and_async_stream_emit_completed_usage(
+ tmp_path: Path,
+):
+ requests: list[httpx.Request] = []
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ requests.append(request)
+ return httpx.Response(
+ 200,
+ headers={"Content-Type": "text/event-stream"},
+ content=_responses_sse(),
+ )
+
+ llm = GrokOAuth(
+ oauth_credential_path=str(tmp_path / "auth.json"),
+ oauth_access_token="stream-access",
+ oauth_refresh_token="stream-refresh",
+ oauth_expires_at_ms=int(time.time() * 1000) + 3_600_000,
+ http_client=httpx.Client(transport=httpx.MockTransport(handler)),
+ async_http_client=httpx.AsyncClient(transport=httpx.MockTransport(handler)),
+ )
+
+ async def collect_async(): # type: ignore[no-untyped-def]
+ stream = await llm.astream_chat(
+ [ChatMessage(role=MessageRole.USER, content="hello")],
+ temperature=0.3,
+ top_p=0.6,
+ reasoning={"effort": "low"},
+ store=True,
+ )
+ return [response async for response in stream]
+
+ try:
+ sync_responses = list(
+ llm.stream_chat(
+ [ChatMessage(role=MessageRole.USER, content="hello")],
+ temperature=0.3,
+ top_p=0.6,
+ reasoning={"effort": "low"},
+ store=True,
+ )
+ )
+ async_responses = asyncio.run(collect_async())
+ finally:
+ llm._client.close()
+ asyncio.run(llm._aclient.close())
+
+ for responses in (sync_responses, async_responses):
+ assert [response.delta for response in responses] == ["hel", ""]
+ assert responses[-1].message.content == "hello"
+ assert responses[-1].raw.type == "response.completed"
+ assert responses[-1].additional_kwargs["usage"].total_tokens == 11
+ usage = get_usage_from_response("GrokOAuth", responses[-1])
+ assert (usage.request_tokens, usage.response_tokens, usage.total_tokens) == (
+ 8,
+ 3,
+ 11,
+ )
+
+ assert len(requests) == 2
+ for request in requests:
+ payload = json.loads(request.content)
+ assert payload["stream"] is True
+ assert payload["store"] is False
+ assert request.headers["Authorization"] == "Bearer stream-access"
+ assert request.headers["X-XAI-Token-Auth"] == "xai-grok-cli"
+ assert request.headers["x-grok-model-override"] == DEFAULT_GROK_MODEL
+ assert (
+ request.headers[GROK_CLI_COMPAT_VERSION_HEADER] == GROK_CLI_COMPAT_VERSION
+ )
+ assert {"temperature", "top_p", "reasoning"}.isdisjoint(payload)
+
+
+def test_oauth_adapter_async_responses_request_uses_proxy_auth_headers(
+ tmp_path: Path,
+):
+ requests: list[httpx.Request] = []
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ requests.append(request)
+ return httpx.Response(
+ 200,
+ json={
+ "id": "resp_test",
+ "created_at": int(time.time()),
+ "model": DEFAULT_GROK_MODEL,
+ "object": "response",
+ "output": [],
+ "parallel_tool_calls": True,
+ "tool_choice": "auto",
+ "tools": [],
+ "status": "completed",
+ },
+ )
+
+ async def run() -> str:
+ async_http_client = httpx.AsyncClient(transport=httpx.MockTransport(handler))
+ llm = GrokOAuth(
+ oauth_credential_path=str(tmp_path / "auth.json"),
+ oauth_access_token="adapter-access",
+ oauth_refresh_token="adapter-refresh",
+ oauth_expires_at_ms=int(time.time() * 1000) + 3_600_000,
+ async_http_client=async_http_client,
+ )
+ try:
+ response = await llm._aclient.responses.create(
+ model=DEFAULT_GROK_MODEL,
+ input="hello",
+ store=False,
+ )
+ return response.id
+ finally:
+ await llm._aclient.close()
+
+ assert asyncio.run(run()) == "resp_test"
+ assert len(requests) == 1
+ request = requests[0]
+ assert str(request.url) == f"{DEFAULT_GROK_OAUTH_PROXY}/responses"
+ assert request.headers["Authorization"] == "Bearer adapter-access"
+ assert request.headers["X-XAI-Token-Auth"] == "xai-grok-cli"
+ assert request.headers["x-grok-model-override"] == DEFAULT_GROK_MODEL
+ assert request.headers[GROK_CLI_COMPAT_VERSION_HEADER] == GROK_CLI_COMPAT_VERSION
+
+
+def test_oauth_responses_adapter_pins_proxy_and_omits_controls(tmp_path: Path):
+ llm = GrokOAuth(
+ oauth_credential_path=str(tmp_path / "auth.json"),
+ oauth_access_token="access",
+ oauth_refresh_token="refresh",
+ oauth_expires_at_ms=int(time.time() * 1000) + 3_600_000,
+ store=True,
+ track_previous_responses=True,
+ reasoning_options={"effort": "low"},
+ context_window=123,
+ default_headers={GROK_CLI_COMPAT_VERSION_HEADER: "caller-version"},
+ additional_kwargs={"presence_penalty": 0.5, "stop": ["done"]},
+ )
+
+ model_kwargs = llm._get_model_kwargs(
+ model="caller-selected-model",
+ top_p=0.5,
+ temperature=0.2,
+ frequency_penalty=0.3,
+ reasoning={"effort": "high"},
+ extra_body={
+ "model": "extra-body-model",
+ "store": True,
+ "temperature": 0.9,
+ "top_p": 0.8,
+ "presence_penalty": 0.7,
+ "frequency_penalty": 0.6,
+ "stop": "extra-stop",
+ "reasoning": {"effort": "low"},
+ "metadata": {"safe": "value"},
+ },
+ )
+ assert llm.api_base == DEFAULT_GROK_OAUTH_PROXY
+ assert llm.metadata.context_window == DEFAULT_GROK_CONTEXT_WINDOW
+ assert llm.metadata.is_function_calling_model
+ assert llm._tokenizer is None
+ assert (
+ llm.default_headers[GROK_CLI_COMPAT_VERSION_HEADER] == GROK_CLI_COMPAT_VERSION
+ )
+ assert model_kwargs["store"] is False
+ assert model_kwargs["model"] == DEFAULT_GROK_MODEL
+ assert model_kwargs["extra_body"] == {"metadata": {"safe": "value"}}
+ for key in (
+ "temperature",
+ "top_p",
+ "presence_penalty",
+ "frequency_penalty",
+ "stop",
+ "reasoning",
+ ):
+ assert key not in model_kwargs
+
+ with pytest.raises(ValueError, match="not supported with XAI OAuth"):
+ GrokOAuth(
+ model="other-model",
+ oauth_credential_path=str(tmp_path / "other.json"),
+ )
+
+ no_fallback = GrokOAuth(
+ api_key="xai-api-key-must-be-ignored",
+ oauth_credential_path=str(tmp_path / "missing.json"),
+ )
+ assert no_fallback.api_key == "oauth"
+ with pytest.raises(ValueError, match="No XAI OAuth credentials"):
+ no_fallback._oauth_manager.get_valid_credentials()
+
+
+def test_structured_predict_sanitizes_runtime_kwargs(monkeypatch, tmp_path: Path):
+ llm = GrokOAuth(
+ oauth_credential_path=str(tmp_path / "auth.json"),
+ oauth_access_token="access",
+ oauth_refresh_token="refresh",
+ oauth_expires_at_ms=int(time.time() * 1000) + 3_600_000,
+ )
+ captured: dict[str, object] = {}
+
+ def fake_parse(**kwargs): # type: ignore[no-untyped-def]
+ captured.update(kwargs)
+ return SimpleNamespace(output_parsed="parsed")
+
+ monkeypatch.setattr(
+ "mobilerun.agent.utils.oauth.grok_oauth_llm.to_openai_message_dicts",
+ lambda messages, **kwargs: [{"role": "user", "content": "formatted"}],
+ )
+ llm._client = SimpleNamespace(responses=SimpleNamespace(parse=fake_parse))
+ prompt = SimpleNamespace(format_messages=lambda **kwargs: ["message"])
+ result = llm.structured_predict(
+ object,
+ prompt,
+ llm_kwargs={
+ "temperature": 0.2,
+ "top_p": 0.5,
+ "presence_penalty": 0.1,
+ "reasoning": {"effort": "low"},
+ "store": True,
+ "tool_choice": "none",
+ "metadata": {"safe": "value"},
+ "model": "caller-selected-model",
+ "extra_body": {
+ "model": "extra-body-model",
+ "store": True,
+ "tool_choice": "required",
+ "reasoning": {"effort": "low"},
+ "temperature": 0.9,
+ "metadata": {"nested": "safe"},
+ },
+ },
+ )
+
+ assert result == "parsed"
+ assert captured == {
+ "model": DEFAULT_GROK_MODEL,
+ "input": [{"role": "user", "content": "formatted"}],
+ "text_format": object,
+ "store": False,
+ "metadata": {"safe": "value"},
+ "extra_body": {"metadata": {"nested": "safe"}},
+ }
+ assert "tool_choice" not in captured
+ assert llm.store is False
+
+
+def test_async_structured_predict_sanitizes_runtime_kwargs(monkeypatch, tmp_path: Path):
+ llm = GrokOAuth(
+ oauth_credential_path=str(tmp_path / "auth.json"),
+ oauth_access_token="access",
+ oauth_refresh_token="refresh",
+ oauth_expires_at_ms=int(time.time() * 1000) + 3_600_000,
+ )
+ captured: dict[str, object] = {}
+
+ async def fake_parse(**kwargs): # type: ignore[no-untyped-def]
+ captured.update(kwargs)
+ return SimpleNamespace(output_parsed="parsed")
+
+ monkeypatch.setattr(
+ "mobilerun.agent.utils.oauth.grok_oauth_llm.to_openai_message_dicts",
+ lambda messages, **kwargs: [{"role": "user", "content": "formatted"}],
+ )
+ llm._aclient = SimpleNamespace(responses=SimpleNamespace(parse=fake_parse))
+ prompt = SimpleNamespace(format_messages=lambda **kwargs: ["message"])
+ result = asyncio.run(
+ llm.astructured_predict(
+ object,
+ prompt,
+ llm_kwargs={
+ "temperature": 0.2,
+ "top_p": 0.5,
+ "reasoning": {"effort": "low"},
+ "store": True,
+ "tool_choice": "required",
+ "metadata": {"safe": "value"},
+ "model": "caller-selected-model",
+ "extra_body": {
+ "model": "extra-body-model",
+ "store": True,
+ "tool_choice": "required",
+ "reasoning": {"effort": "low"},
+ "top_p": 0.8,
+ "metadata": {"nested": "safe"},
+ },
+ },
+ )
+ )
+
+ assert result == "parsed"
+ assert captured == {
+ "model": DEFAULT_GROK_MODEL,
+ "input": [{"role": "user", "content": "formatted"}],
+ "text_format": object,
+ "store": False,
+ "metadata": {"safe": "value"},
+ "extra_body": {"metadata": {"nested": "safe"}},
+ }
+ assert "tool_choice" not in captured
+
+
+def test_grok_integration_source_does_not_bridge_to_external_credentials():
+ source = Path("mobilerun/agent/utils/oauth/grok_oauth_llm.py").read_text(
+ encoding="utf-8"
+ )
+ forbidden = (
+ "." + "grok" + "/" + "auth.json",
+ "GROK" + "_HOME",
+ "subprocess" + ".run",
+ )
+ assert not any(value in source for value in forbidden)
diff --git a/tests/test_oauth_login_timeout.py b/tests/test_oauth_login_timeout.py
new file mode 100644
index 00000000..041c05a5
--- /dev/null
+++ b/tests/test_oauth_login_timeout.py
@@ -0,0 +1,263 @@
+from __future__ import annotations
+
+import math
+from types import SimpleNamespace
+
+import httpx
+import pytest
+
+from mobilerun.agent.utils.oauth import (
+ anthropic_oauth_llm,
+ gemini_oauth_code_assist_llm,
+ openai_oauth_llm,
+)
+from mobilerun.agent.utils.oauth.anthropic_oauth_llm import AnthropicOAuthLLM
+from mobilerun.agent.utils.oauth.gemini_oauth_code_assist_llm import (
+ GeminiOAuthCodeAssistLLM,
+)
+from mobilerun.agent.utils.oauth.login_timeout import OAuthLoginDeadline
+from mobilerun.agent.utils.oauth.openai_oauth_llm import OpenAIOAuth
+from mobilerun.config_manager.auth_profile_store import AuthProfileStore
+
+
+class _FakeClock:
+ def __init__(self) -> None:
+ self.now = 10.0
+ self.sleeps: list[float] = []
+
+ def __call__(self) -> float:
+ return self.now
+
+ def sleep(self, seconds: float) -> None:
+ self.sleeps.append(seconds)
+ self.now += seconds
+
+
+@pytest.mark.parametrize("timeout", [0, -1, math.nan, math.inf, -math.inf])
+def test_oauth_login_deadline_rejects_non_positive_or_non_finite_timeout(
+ timeout: float,
+) -> None:
+ with pytest.raises(ValueError, match="finite positive"):
+ OAuthLoginDeadline(timeout)
+
+
+def test_oauth_login_deadline_caps_requests_and_never_resets() -> None:
+ clock = _FakeClock()
+ deadline = OAuthLoginDeadline(10, clock=clock, sleeper=clock.sleep)
+
+ assert deadline.remaining(cap=3) == 3
+ clock.now += 4
+ assert deadline.remaining() == 6
+
+ limited = deadline.limited_to(2)
+ limited.sleep(1)
+ assert clock.sleeps == [1]
+ assert limited.remaining() == 1
+
+ with pytest.raises(TimeoutError, match="OAuth login timed out"):
+ limited.sleep(2)
+ assert clock.sleeps == [1, 1]
+ assert deadline.remaining() == 4
+
+
+def test_auth_profile_commit_honors_deadline_and_lock_timeout(tmp_path) -> None:
+ path = tmp_path / "auth-profiles.json"
+ store = AuthProfileStore(path)
+
+ with pytest.raises(TimeoutError, match="expired before commit"):
+ store.update_slot(
+ "slot",
+ {"accessToken": "secret"},
+ before_commit=lambda: (_ for _ in ()).throw(
+ TimeoutError("expired before commit")
+ ),
+ )
+ assert not path.exists()
+
+ with store.transaction():
+ with pytest.raises(TimeoutError):
+ store.update_slot(
+ "slot",
+ {"accessToken": "secret"},
+ lock_timeout=0.01,
+ )
+ assert not path.exists()
+
+
+@pytest.mark.parametrize("open_browser", [False, True])
+def test_openai_device_flow_preserves_deadline_and_browser_preference(
+ monkeypatch,
+ tmp_path,
+ open_browser: bool,
+) -> None:
+ clock = _FakeClock()
+ deadline = OAuthLoginDeadline(20, clock=clock, sleeper=clock.sleep)
+ opened: list[str] = []
+ requests: list[tuple[str, float]] = []
+ responses = iter(
+ [
+ httpx.Response(503),
+ httpx.Response(
+ 200,
+ json={
+ "device_auth_id": "device",
+ "user_code": "ABCD",
+ "verification_uri": "https://auth.openai.test/device",
+ "expires_in": 60,
+ "interval": 1,
+ },
+ ),
+ httpx.Response(403),
+ httpx.Response(
+ 200,
+ json={"authorization_code": "code", "code_verifier": "verifier"},
+ ),
+ httpx.Response(
+ 200,
+ json={
+ "access_token": "access",
+ "refresh_token": "refresh",
+ "expires_in": 3600,
+ },
+ ),
+ ]
+ )
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ requests.append(
+ (request.url.path, float(request.extensions["timeout"]["read"]))
+ )
+ clock.now += 0.5
+ return next(responses)
+
+ llm = OpenAIOAuth(
+ model="gpt-5.5",
+ oauth_credential_path=str(tmp_path / "auth.json"),
+ )
+ llm._oauth_manager.http_client = httpx.Client(
+ transport=httpx.MockTransport(handler)
+ )
+ monkeypatch.setattr(
+ openai_oauth_llm,
+ "open_browser_async",
+ lambda url, opener: opened.append(url),
+ )
+
+ credentials = llm._login_device_code(
+ open_browser=open_browser,
+ deadline=deadline,
+ )
+
+ assert credentials.access_token == "access"
+ assert llm._oauth_manager.credential_store.load() == credentials
+ assert [path for path, _ in requests] == [
+ "/api/accounts/deviceauth/usercode",
+ "/api/accounts/deviceauth/usercode",
+ "/api/accounts/deviceauth/token",
+ "/api/accounts/deviceauth/token",
+ "/oauth/token",
+ ]
+ assert [timeout for _, timeout in requests] == pytest.approx(
+ [20.0, 17.5, 17.0, 15.5, 15.0]
+ )
+ assert clock.sleeps == [2, 1]
+ assert opened == (["https://auth.openai.test/device"] if open_browser else [])
+
+
+@pytest.mark.parametrize(
+ ("llm", "exchange_kwargs"),
+ (
+ (
+ AnthropicOAuthLLM(credential_path=None, timeout=30),
+ {
+ "code": "code",
+ "redirect_uri": "https://example.test/callback",
+ "code_verifier": "verifier",
+ "state": "state",
+ },
+ ),
+ (
+ GeminiOAuthCodeAssistLLM(credential_path=None, timeout=30),
+ {
+ "code": "code",
+ "redirect_uri": "https://example.test/callback",
+ "code_verifier": "verifier",
+ },
+ ),
+ ),
+)
+def test_anthropic_and_gemini_exchange_use_one_remaining_budget(
+ llm,
+ exchange_kwargs,
+) -> None:
+ clock = _FakeClock()
+ deadline = OAuthLoginDeadline(5, clock=clock, sleeper=clock.sleep)
+ timeouts: list[float] = []
+
+ def post(*args, **kwargs): # type: ignore[no-untyped-def]
+ timeouts.append(kwargs["timeout"])
+ clock.now += 1
+ return SimpleNamespace(
+ raise_for_status=lambda: None,
+ json=lambda: {
+ "access_token": "access",
+ "refresh_token": "refresh",
+ "expires_in": 3600,
+ },
+ )
+
+ llm._session = SimpleNamespace(post=post)
+
+ assert (
+ llm._exchange_authorization_code(
+ **exchange_kwargs,
+ deadline=deadline,
+ )
+ == "access"
+ )
+ assert timeouts == [5]
+
+
+@pytest.mark.parametrize("open_browser", [False, True])
+@pytest.mark.parametrize(
+ ("module", "llm_class", "constructor_kwargs"),
+ (
+ (anthropic_oauth_llm, AnthropicOAuthLLM, {"credential_path": None}),
+ (
+ gemini_oauth_code_assist_llm,
+ GeminiOAuthCodeAssistLLM,
+ {"credential_path": None},
+ ),
+ ),
+)
+def test_anthropic_and_gemini_bind_fallback_preserves_deadline_and_browser(
+ monkeypatch,
+ module,
+ llm_class,
+ constructor_kwargs,
+ open_browser: bool,
+) -> None:
+ deadline = OAuthLoginDeadline(5)
+ calls: list[dict[str, object]] = []
+ monkeypatch.setattr(module, "_is_headless_environment", lambda: False)
+ monkeypatch.setattr(
+ module,
+ "HTTPServer",
+ lambda *args, **kwargs: (_ for _ in ()).throw(OSError("unavailable")),
+ )
+ monkeypatch.setattr(
+ llm_class,
+ "login_headless",
+ lambda self, **kwargs: (calls.append(kwargs), "access")[1],
+ )
+
+ assert (
+ llm_class(**constructor_kwargs).login(
+ open_browser=open_browser,
+ deadline=deadline,
+ )
+ == "access"
+ )
+
+ assert calls[0]["deadline"] is deadline
+ assert calls[0]["open_browser"] is open_browser
diff --git a/tests/test_usage.py b/tests/test_usage.py
index 6e7e6bcd..c8683b25 100644
--- a/tests/test_usage.py
+++ b/tests/test_usage.py
@@ -24,6 +24,14 @@ def _openai_responses_chat_response() -> ChatResponse:
)
+def _streamed_responses_chat_response(*, raw, additional_kwargs=None) -> ChatResponse:
+ return ChatResponse(
+ message=ChatMessage(role=MessageRole.ASSISTANT, content="ok"),
+ raw=raw,
+ additional_kwargs=additional_kwargs or {},
+ )
+
+
def test_track_usage_supports_mobilerun_openai_responses_wrapper() -> None:
llm = load_llm("OpenAIResponses", model="gpt-5.5", api_key="stub")
@@ -64,6 +72,87 @@ def test_openai_responses_class_name_extracts_usage_from_response() -> None:
assert usage.requests == 1
+def test_openai_responses_extracts_usage_from_completed_stream_object() -> None:
+ chat_response = _streamed_responses_chat_response(
+ raw=SimpleNamespace(
+ type="response.completed",
+ response=SimpleNamespace(
+ usage=SimpleNamespace(
+ input_tokens=7,
+ output_tokens=5,
+ total_tokens=12,
+ )
+ ),
+ )
+ )
+
+ usage = get_usage_from_response("OpenAIResponses", chat_response)
+
+ assert usage.request_tokens == 7
+ assert usage.response_tokens == 5
+ assert usage.total_tokens == 12
+ assert usage.requests == 1
+
+
+def test_openai_responses_extracts_usage_from_completed_stream_dict() -> None:
+ chat_response = _streamed_responses_chat_response(
+ raw={
+ "type": "response.completed",
+ "response": {
+ "usage": {
+ "input_tokens": 11,
+ "output_tokens": 4,
+ "total_tokens": 15,
+ }
+ },
+ }
+ )
+
+ usage = get_usage_from_response("OpenAIResponses", chat_response)
+
+ assert usage.request_tokens == 11
+ assert usage.response_tokens == 4
+ assert usage.total_tokens == 15
+
+
+def test_openai_responses_extracts_completed_stream_additional_usage_fallback() -> None:
+ chat_response = _streamed_responses_chat_response(
+ raw={"type": "response.completed"},
+ additional_kwargs={
+ "usage": {
+ "input_tokens": 13,
+ "output_tokens": 6,
+ "total_tokens": 19,
+ }
+ },
+ )
+
+ usage = get_usage_from_response("xai_oauth", chat_response)
+
+ assert usage.request_tokens == 13
+ assert usage.response_tokens == 6
+ assert usage.total_tokens == 19
+
+
+def test_openai_responses_extracts_object_usage_without_raw_event() -> None:
+ chat_response = _streamed_responses_chat_response(
+ raw=None,
+ additional_kwargs={
+ "usage": SimpleNamespace(
+ input_tokens=17,
+ output_tokens=8,
+ total_tokens=25,
+ )
+ },
+ )
+
+ usage = get_usage_from_response("OpenAIResponses", chat_response)
+
+ assert usage.request_tokens == 17
+ assert usage.response_tokens == 8
+ assert usage.total_tokens == 25
+
+
def test_track_usage_supports_mobilerun_anthropic_wrapper() -> None:
llm = load_llm("Anthropic", model="claude-opus-4-8", api_key="stub")
diff --git a/uv.lock b/uv.lock
index cf15395a..e7ad0231 100644
--- a/uv.lock
+++ b/uv.lock
@@ -1685,6 +1685,7 @@ dependencies = [
{ name = "aiofiles" },
{ name = "arize-phoenix" },
{ name = "async-adbutils" },
+ { name = "filelock" },
{ name = "httpx" },
{ name = "inquirerpy" },
{ name = "llama-index" },
@@ -1700,6 +1701,7 @@ dependencies = [
{ name = "mobilerun-sdk" },
{ name = "posthog" },
{ name = "pydantic" },
+ { name = "pyjwt", extra = ["crypto"] },
{ name = "rich" },
]
@@ -1745,6 +1747,7 @@ requires-dist = [
{ name = "async-adbutils" },
{ name = "bandit", marker = "extra == 'dev'", specifier = ">=1.8.6" },
{ name = "black", marker = "extra == 'dev'", specifier = "==25.9.0" },
+ { name = "filelock", specifier = ">=3.20.0" },
{ name = "httpx", specifier = ">=0.27.0" },
{ name = "inquirerpy", specifier = ">=0.3.4" },
{ name = "langfuse", marker = "extra == 'langfuse'", specifier = "==3.12.1" },
@@ -1766,6 +1769,7 @@ requires-dist = [
{ name = "openinference-instrumentation-llama-index", marker = "extra == 'langfuse'", specifier = ">=3.0.0" },
{ name = "posthog", specifier = ">=6.7.6" },
{ name = "pydantic", specifier = ">=2.11.10" },
+ { name = "pyjwt", extras = ["crypto"], specifier = ">=2.10.0" },
{ name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0.0" },
{ name = "rich", specifier = ">=14.1.0" },
{ name = "ruff", marker = "extra == 'dev'", specifier = ">=0.13.0" },