From b0c4b7cc7ee5aa2af7f99cce4c85366f473d891e Mon Sep 17 00:00:00 2001 From: "rasul.osmanbayli" Date: Wed, 12 Aug 2026 17:21:03 +0400 Subject: [PATCH 1/4] Add Grok API and OAuth support --- README.md | 19 +- docs/guides/cli.mdx | 48 + docs/quickstart.mdx | 10 +- docs/sdk/configuration.mdx | 21 +- mobilerun/agent/providers/grok.py | 69 + mobilerun/agent/providers/registry.py | 33 + mobilerun/agent/providers/setup_service.py | 8 +- mobilerun/agent/usage.py | 30 +- mobilerun/agent/utils/llm_picker.py | 143 +- .../agent/utils/oauth/anthropic_oauth_llm.py | 24 +- .../oauth/gemini_oauth_code_assist_llm.py | 24 +- mobilerun/agent/utils/oauth/grok_oauth_llm.py | 949 ++++++++++ .../agent/utils/oauth/openai_oauth_llm.py | 28 +- mobilerun/cli/configure_wizard.py | 16 +- mobilerun/cli/main.py | 68 +- mobilerun/cli/oauth_actions.py | 47 +- mobilerun/config_example.yaml | 18 + .../config_manager/auth_profile_store.py | 160 ++ mobilerun/config_manager/credential_paths.py | 1 + mobilerun/config_manager/env_keys.py | 54 +- pyproject.toml | 2 + tests/e2e/test_grok_android16.py | 1554 +++++++++++++++++ tests/test_grok_android_e2e_helpers.py | 174 ++ tests/test_grok_api.py | 553 ++++++ tests/test_grok_cli.py | 178 ++ tests/test_grok_oauth.py | 1496 ++++++++++++++++ tests/test_llm_picker.py | 1 + tests/test_usage.py | 89 + uv.lock | 4 + 29 files changed, 5679 insertions(+), 142 deletions(-) create mode 100644 mobilerun/agent/providers/grok.py create mode 100644 mobilerun/agent/utils/oauth/grok_oauth_llm.py create mode 100644 mobilerun/config_manager/auth_profile_store.py create mode 100644 tests/e2e/test_grok_android16.py create mode 100644 tests/test_grok_android_e2e_helpers.py create mode 100644 tests/test_grok_api.py create mode 100644 tests/test_grok_cli.py create mode 100644 tests/test_grok_oauth.py diff --git a/README.md b/README.md index c4204f01..12b3ee23 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, Grok, 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,22 @@ 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`. + +Grok supports either an xAI API key or a native xAI OAuth login managed by +Mobilerun: + +```bash +# Direct xAI API +export XAI_API_KEY=your-key +mobilerun configure --provider grok --auth-mode api_key --model grok-4.5 + +# Native Grok OAuth (use --device-code on SSH/headless machines) +mobilerun configure grok +``` + +Mobilerun stores this login in its own `grokOauth` credential slot. It does not +invoke the Grok CLI or read `~/.grok/auth.json`. ### 4. Run your first command diff --git a/docs/guides/cli.mdx b/docs/guides/cli.mdx index 04a019a7..b0dff4d8 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,8 @@ 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` | +| grok_oauth | Included by default | Mobilerun-managed xAI OAuth credentials | | OpenAILike | Included by default | Varies by provider | | OpenRouter | Included by default | `OPENROUTER_API_KEY` | | Ollama | Included by default | None (local) | @@ -161,6 +169,45 @@ mobilerun run "Enable 2FA" \ | DeepSeek | Included by default | `DEEPSEEK_API_KEY` | | MiniMax | Included by default | `MINIMAX_API_KEY` | +### Grok API key and OAuth + +Grok 4.5 is available through the public xAI API and through a Grok account. +For public API billing, set `XAI_API_KEY` or paste the key in the configure +wizard: + +```bash +mobilerun configure \ + --provider grok \ + --auth-mode api_key \ + --model grok-4.5 +``` + +For Grok OAuth, let Mobilerun run its native xAI login flow and save the result +in the shared `auth-profiles.json` store: + +```bash +# Browser login +mobilerun configure grok + +# Device-code login for SSH/headless hosts +mobilerun configure grok --device-code --no-browser +``` + +You can also choose **Grok** and **oauth** in `mobilerun configure`. The +resulting profile uses provider `grok_oauth`; API-key profiles use provider +`XAI`. OAuth credentials are stored under `grokOauth` with owner-only file +permissions. This credential is separate from `~/.grok/auth.json`: Mobilerun +does not invoke the Grok CLI, read its credential file, or treat an OAuth +failure as permission to fall back to `XAI_API_KEY`. + + + Grok OAuth currently interoperates with xAI's public Grok Build OAuth client + and subscription proxy contract. xAI does not document that contract as a + stable third-party integration, so xAI may require a fresh login or a + compatibility update in a future release. Use `XAI` with an API key when you + need the documented public API boundary. + + ### MiniMax endpoints and credentials Run `mobilerun configure`, choose MiniMax, and then select the API region where @@ -551,6 +598,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..3a9117f0 100644 --- a/docs/quickstart.mdx +++ b/docs/quickstart.mdx @@ -104,8 +104,16 @@ export OPENAI_API_KEY=your-api-key-here # For Anthropic Claude export ANTHROPIC_API_KEY=your-api-key-here + +# For xAI Grok +export XAI_API_KEY=your-api-key-here ``` +To use your Grok account instead of an API key, run +`mobilerun configure grok` (add `--device-code --no-browser` on a headless +host). Mobilerun runs its own xAI OAuth flow and never invokes the Grok CLI or +reads `~/.grok/auth.json`. + ### Run Your First Command via CLI Now you're ready to control your device with natural language: @@ -125,7 +133,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, grok_oauth, 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..8774cb47 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, grok_oauth, Ollama, Anthropic, GoogleGenAI, DeepSeek, MiniMax) - `--model MODEL` - LLM model name - `--temperature FLOAT` - LLM temperature - `--steps INT` - Max steps @@ -635,11 +635,30 @@ 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 ``` +Grok 4.5 supports both the public xAI API and Grok OAuth. Use the `XAI` +runtime provider for API-key calls: + +```bash +mobilerun configure \ + --provider grok \ + --auth-mode api_key \ + --model grok-4.5 +``` + +For OAuth, run `mobilerun configure grok` (or add +`--device-code --no-browser` on a headless host). Mobilerun performs its own +xAI OAuth flow, stores the credential under `grokOauth` in the shared platform +`credentials/auth-profiles.json` file, and generates a `grok_oauth` profile. +The Grok CLI is not required or invoked, and its `~/.grok/auth.json` remains +separate. API-key profiles use `apiKeys.xai` / `XAI_API_KEY` and the `XAI` +runtime provider. + MiniMax uses `https://api.minimax.io/v1` for global accounts and `https://api.minimaxi.com/v1` for Mainland China accounts. The configure wizard asks for the matching region and selects `MiniMax-M3` by default: diff --git a/mobilerun/agent/providers/grok.py b/mobilerun/agent/providers/grok.py new file mode 100644 index 00000000..e350a98f --- /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..f9f74cdd 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="grok", + display_name="Grok", + 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="grok_oauth", + runtime_provider_name="grok_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 == "grok": + 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..02989635 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") @@ -150,7 +151,6 @@ def create_profile_for_variant( resolved_model, ) kwargs: dict[str, str | int] = dict(DEFAULT_KWARGS_BY_VARIANT.get(variant.id, {})) - env_slot = VARIANT_ENV_KEY_SLOT.get(variant.id) runtime_provider_name = ( variant.runtime_transport_provider_name or variant.runtime_provider_name ) @@ -176,10 +176,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, ) diff --git a/mobilerun/agent/usage.py b/mobilerun/agent/usage.py index 13714a3b..1cdb3d9d 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", + "grok_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..45bc1039 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,9 @@ "openai_like": "OpenAILike", "zai": "ZAI", "z.ai": "ZAI", + "grok": "XAI", + "xai": "XAI", + "x.ai": "XAI", } ZAI_GLOBAL_API_BASE = "https://api.z.ai/api/paas/v4" @@ -184,12 +195,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 +227,23 @@ 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 +251,31 @@ 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 +286,43 @@ 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 +527,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 == "grok_oauth": + model = normalize_model_id_for_variant("grok", "oauth", model) elif provider_name == "openai_oauth": model = normalize_model_id_for_variant("openai", "oauth", model) kwargs["model"] = model @@ -460,6 +560,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 == "grok_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 +593,37 @@ 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 public runtime aliases (`grok`, `xai`, and `x.ai`) should be + # useful without a separately generated profile as well. Keep their + # implicit model aligned with the first-class provider 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..0cfca70a 100644 --- a/mobilerun/agent/utils/oauth/anthropic_oauth_llm.py +++ b/mobilerun/agent/utils/oauth/anthropic_oauth_llm.py @@ -37,6 +37,7 @@ anthropic_model_omits_sampling_params, strip_anthropic_sampling_params, ) +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 @@ -266,18 +267,7 @@ def _persist_credentials(self) -> 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"] = { + AuthProfileStore(path).update_slot("claudeAiOauth", { "accessToken": self._cached_access_token, "refreshToken": self._cached_refresh_token, "expiresAt": ( @@ -286,15 +276,7 @@ def _persist_credentials(self) -> None: 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 + }) def _token_headers(self) -> Dict[str, str]: headers = { 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..f5efd952 100644 --- a/mobilerun/agent/utils/oauth/gemini_oauth_code_assist_llm.py +++ b/mobilerun/agent/utils/oauth/gemini_oauth_code_assist_llm.py @@ -31,6 +31,7 @@ from llama_index.core.llms.callbacks import llm_chat_callback, llm_completion_callback from llama_index.core.llms.custom import CustomLLM +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" @@ -296,18 +297,7 @@ def _persist_credentials(self) -> None: 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] = { + AuthProfileStore(path).update_slot(self.credential_slot, { "access_token": self._cached_access_token, "refresh_token": self._cached_refresh_token, "token_type": "Bearer", @@ -316,15 +306,7 @@ def _persist_credentials(self) -> None: 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 + }) def _metadata_payload(self) -> Dict[str, str]: return { 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..218f4dc9 --- /dev/null +++ b/mobilerun/agent/utils/oauth/grok_oauth_llm.py @@ -0,0 +1,949 @@ +"""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, 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.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("Grok OAuth profile has an unexpected credential type.") + access_token = payload.get("accessToken") + if not isinstance(access_token, str) or not access_token: + raise ValueError("Grok OAuth profile is missing accessToken.") + issuer = payload.get("issuer") + client_id = payload.get("clientId") + if issuer != DEFAULT_GROK_OAUTH_ISSUER: + raise ValueError("Grok OAuth profile has an unexpected issuer.") + if client_id != DEFAULT_GROK_OAUTH_CLIENT_ID: + raise ValueError("Grok 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("Grok 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) -> None: + self.profile_store.update_slot( + DEFAULT_GROK_OAUTH_SLOT, credentials.to_payload() + ) + + +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 + ) + + def validate(self, token: str, *, nonce: str | None) -> dict[str, Any]: + signing_key = self._jwks_client.get_signing_key_from_jwt(token).key + 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.") + 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, + ) -> GrokOAuthCredentials: + 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: + self.id_token_validator.validate(id_token, nonce=nonce) + + 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.") + return 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), + ) + + def _post_form( + self, + url: str, + *, + data: dict[str, str], + headers: dict[str, str], + context: str, + retry_transient: bool, + ) -> 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): + try: + response = self.http_client.post( + url, + headers=headers, + data=data, + timeout=self.request_timeout, + ) + except (httpx.ConnectError, httpx.TimeoutException) as exc: + if attempt < len(backoffs): + self.sleep(backoffs[attempt]) + continue + raise GrokOAuthError( + f"{context} failed due to a transient network error." + ) from exc + if 500 <= response.status_code < 600 and attempt < len(backoffs): + 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, + ) -> 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, + ) + 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( + "Grok OAuth refresh was rejected; re-login is required." + ) + raise GrokOAuthError( + f"xAI token request failed ({error or response.status_code})." + ) + return _safe_json_object(response, context="xAI token response") + + def set_initial_credentials(self, credentials: GrokOAuthCredentials) -> None: + with self._thread_lock: + self.credential_store.save(credentials) + self._credentials = credentials + + def exchange_authorization_code( + self, *, code: str, redirect_uri: str, code_verifier: str, nonce: str + ) -> 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, + } + ) + credentials = self._credentials_from_token_response( + payload, nonce=nonce + ) + self.set_initial_credentials(credentials) + return credentials + + def _refresh(self, credentials: GrokOAuthCredentials) -> GrokOAuthCredentials: + if not credentials.refresh_token: + raise ValueError( + "No Grok OAuth refresh token is available. Run MobileRun's Grok login." + ) + 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 Grok OAuth credentials found. Run MobileRun's Grok login." + ) + + 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 Grok 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, + ) -> GrokOAuthCredentials: + if callback_host != DEFAULT_GROK_OAUTH_CALLBACK_HOST: + raise ValueError("Grok OAuth callback_host must be 127.0.0.1.") + if callback_port != 0: + raise ValueError("Grok OAuth callback_port must be OS-assigned (0).") + if callback_path != DEFAULT_GROK_OAUTH_CALLBACK_PATH: + raise ValueError("Grok OAuth callback_path must be /callback.") + if device_code or _is_headless_environment(): + return self._login_device_code(timeout_seconds=timeout_seconds) + + 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() + + 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 + 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"]) + 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 Grok login complete. You may close this tab." + if ok + else b"MobileRun Grok login failed. Return to the terminal." + ) + done.set() + + 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(timeout_seconds=timeout_seconds) + + 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) + thread.start() + try: + print(f"Open this URL to sign in to xAI:\n{authorization_url}\n") + if open_browser: + webbrowser.open(authorization_url) + if not done.wait(timeout=max(0.0, timeout_seconds)): + raise TimeoutError("Grok OAuth login timed out waiting for callback.") + 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("Grok OAuth callback state mismatch.") + if not result["code"]: + raise GrokOAuthError("Grok 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, + ) + finally: + server.shutdown() + server.server_close() + + def _login_device_code( + self, *, timeout_seconds: float = 1800.0 + ) -> GrokOAuthCredentials: + manager = self._oauth_manager + response = manager.http_client.post( + 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), + }, + timeout=manager.request_timeout, + ) + 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 + deadline = time.monotonic() + min(max(0.0, timeout_seconds), 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" + ) + + while time.monotonic() < deadline: + 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, + ) + if token_response.status_code < 400: + token_payload = _safe_json_object( + token_response, context="xAI device token response" + ) + credentials = manager._credentials_from_token_response( + token_payload + ) + manager.set_initial_credentials(credentials) + 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})." + ) + manager.sleep(min(interval, max(0.0, deadline - time.monotonic()))) + raise TimeoutError("Grok OAuth device authorization timed out.") + + +# Descriptive alias for callers that prefer the full class name. +GrokOAuthLLM = GrokOAuth diff --git a/mobilerun/agent/utils/oauth/openai_oauth_llm.py b/mobilerun/agent/utils/oauth/openai_oauth_llm.py index 70616b1d..1c89b8b1 100644 --- a/mobilerun/agent/utils/oauth/openai_oauth_llm.py +++ b/mobilerun/agent/utils/oauth/openai_oauth_llm.py @@ -41,6 +41,7 @@ from llama_index.llms.openai.utils import to_openai_message_dicts from mobilerun.agent.providers.registry import normalize_model_id_for_variant +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" @@ -273,6 +274,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 +282,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: @@ -289,29 +291,7 @@ def load(self) -> Optional[OpenAIOAuthCredentials]: 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 + self._store.update_slot(self._NESTED_KEY, credentials.to_dict()) class OpenAIOAuthSessionManager: diff --git a/mobilerun/cli/configure_wizard.py b/mobilerun/cli/configure_wizard.py index a6bd8f73..f11c0dca 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 @@ -242,6 +243,7 @@ def _prompt_base_url_for_variant(variant: Any) -> str: "gemini_oauth_code_assist": "geminiAntigravityOauth", "openai_oauth": "openaiOauth", "anthropic_oauth": "claudeAiOauth", + "grok_oauth": "grokOauth", } @@ -321,6 +323,13 @@ def _prepare_variant_auth( callbacks.run_gemini_oauth_login( credential_path=credential_path, model=selected_model ) + elif variant.id == "grok_oauth" and credential_path: + if callbacks.run_grok_oauth_login is None: + raise RuntimeError("Grok 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: @@ -682,10 +691,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. + # A fixed provider should enter that provider's flow immediately. Missing + # auth/model/key values remain interactive, but the supplied provider (and + # auth mode, when present) must not be discarded by the top-level menu. provider_configured = False - if provider_is_fixed and model_is_fixed: + if provider_is_fixed: provider_configured = _configure_provider_model( console, config, diff --git a/mobilerun/cli/main.py b/mobilerun/cli/main.py index 419023da..34ecc1a1 100644 --- a/mobilerun/cli/main.py +++ b/mobilerun/cli/main.py @@ -26,6 +26,8 @@ ping_portal_tcp, setup_portal, ) +from mobilerun_core_local.driver.ios import discover_ios_device, validate_ios_portal_url +from mobilerun_core_local.driver.visual_remote import VISUAL_REMOTE_CONNECTION from rich.console import Console from rich.panel import Panel from rich.text import Text @@ -53,6 +55,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,12 +63,11 @@ 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 from mobilerun.telemetry import print_telemetry_message -from mobilerun_core_local.driver.ios import discover_ios_device, validate_ios_portal_url -from mobilerun_core_local.driver.visual_remote import VISUAL_REMOTE_CONNECTION # Suppress all warnings warnings.filterwarnings("ignore") @@ -422,6 +424,17 @@ def _run_anthropic_oauth_login(credential_path: str, **kwargs) -> None: _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("Grok", credential_path) + + try: _available_agents = list_agents() except Exception: @@ -447,7 +460,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, grok_oauth, Ollama, Anthropic, anthropic_oauth, GoogleGenAI, gemini_oauth_code_assist, DeepSeek)", default=None, ) @click.option( @@ -1012,7 +1025,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, grok, ollama, openai_like, minimax, zai).", ) @click.option( "--auth-mode", @@ -1052,6 +1065,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, @@ -1208,6 +1222,52 @@ def configure_gemini( ) +@configure.command("grok") +@click.option( + "--credential-path", + default=str(GROK_OAUTH_CREDENTIAL_PATH), + show_default=True, + help="Where to store Mobilerun's Grok OAuth credentials.", +) +@click.option( + "--model", default=None, help="Optional model override for later API calls." +) +@click.option( + "--timeout", + type=float, + default=300.0, + show_default=True, + help="Max seconds to wait for xAI OAuth authentication.", +) +@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 Grok's device-code flow for SSH or other headless environments.", +) +def configure_grok( + credential_path: str, + model: str | None, + timeout: float, + open_browser: bool, + device_code: bool, +): + """Log in to Grok with Mobilerun's native xAI OAuth flow.""" + _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..1c43821c 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,10 @@ 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.openai_oauth_llm import ( DEFAULT_OPENAI_OAUTH_CALLBACK_HOST, DEFAULT_OPENAI_OAUTH_CALLBACK_PATH, @@ -21,6 +21,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 @@ -79,6 +80,26 @@ def run_gemini_oauth_login( 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 = False, +) -> None: + """Authenticate directly with xAI and save MobileRun-owned credentials.""" + llm = GrokOAuth( + model=model or DEFAULT_GROK_MODEL, + oauth_credential_path=credential_path, + ) + llm.login( + open_browser=open_browser and not no_browser, + timeout_seconds=timeout, + device_code=device_code, + ) + + def run_anthropic_setup_token_oauth( *, timeout: float = 300.0, @@ -103,26 +124,12 @@ def run_anthropic_setup_token_oauth( 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"] = { + AuthProfileStore(credential_path).update_slot("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 run_anthropic_oauth_setup(credential_path: str) -> None: diff --git a/mobilerun/config_example.yaml b/mobilerun/config_example.yaml index 0058166f..0f1ec3d0 100644 --- a/mobilerun/config_example.yaml +++ b/mobilerun/config_example.yaml @@ -78,6 +78,24 @@ llm_profiles: # kwargs: # optional kwargs, add api_key in kwargs if not already in .env # max_tokens: 8192 # + # Grok API-key example (reads apiKeys.xai / XAI_API_KEY): + # provider: XAI + # provider_family: grok + # auth_mode: api_key + # model: grok-4.5 + # api_key_source: auto + # base_url: https://api.x.ai/v1 + # api_base: https://api.x.ai/v1 + # kwargs: + # context_window: 500000 + # + # Grok OAuth example (first run `mobilerun configure grok`): + # provider: grok_oauth + # provider_family: grok + # auth_mode: oauth + # model: grok-4.5 + # credential_path: /droidrun/credentials/auth-profiles.json + # # Ollama example โ€” max_tokens is translated to Ollama's num_predict, and # context_window controls num_ctx (defaults to 32768; -1 = model maximum, # which preallocates the full KV cache and can spill to CPU): diff --git a/mobilerun/config_manager/auth_profile_store.py b/mobilerun/config_manager/auth_profile_store.py new file mode 100644 index 00000000..72fc6d4f --- /dev/null +++ b/mobilerun/config_manager/auth_profile_store.py @@ -0,0 +1,160 @@ +"""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") + + +class AuthProfileTransaction(AbstractContextManager["AuthProfileTransaction"]): + """A locked read/modify/write transaction over an auth profile file.""" + + def __init__(self, store: "AuthProfileStore") -> None: + self._store = store + self._lock = FileLock(str(store.lock_path)) + self._profile: dict[str, Any] = {} + self._dirty = False + + def __enter__(self) -> "AuthProfileTransaction": + self._store.path.parent.mkdir(parents=True, exist_ok=True) + self._lock.acquire() + 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) + 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) -> AuthProfileTransaction: + return AuthProfileTransaction(self) + + 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]) -> None: + with self.transaction() 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]) -> 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) + os.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()) + 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/e2e/test_grok_android16.py b/tests/e2e/test_grok_android16.py new file mode 100644 index 00000000..9fd4b5db --- /dev/null +++ b/tests/e2e/test_grok_android16.py @@ -0,0 +1,1554 @@ +"""Opt-in Grok API/OAuth E2E coverage on the dedicated Android 16 emulator. + +These tests exercise the natural-language agent through public ``mobilerun +run`` and use the mobile harness's public ``mobilerun_core.Mobilerun`` surface +for harness-owned device actions and evidence. They are excluded from normal +test runs because they consume live xAI credentials and mutate the emulator's +foreground UI. To run them, set both of the following explicitly:: + + MOBILERUN_RUN_ANDROID_E2E=1 + ANDROID_SERIAL=emulator-5558 + +The host must provide AVD ``mobilerun_agent_bench_api36`` with snapshot +``mobilerun_eval_clean_api36``. The session claims only port 5558, launches and +owns that exact emulator, and tears down only ``emulator-5558``. Every scenario +restores the snapshot, runs ``mobilerun setup``, and verifies the Portal before +starting the task. Agent auto-setup remains disabled during the task itself. +""" + +from __future__ import annotations + +import base64 +import hashlib +import io +import json +import os +import re +import shutil +import socket +import subprocess +import sys +import tempfile +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Any, TextIO + +import pytest + +_EXPECTED_SERIAL = "emulator-5558" +_EMULATOR_PORT = 5558 +_AVD_NAME = "mobilerun_agent_bench_api36" +_SNAPSHOT_NAME = "mobilerun_eval_clean_api36" +_BOOT_TIMEOUT_SECONDS = 240.0 +_RUN_LIVE_E2E = os.environ.get("MOBILERUN_RUN_ANDROID_E2E") == "1" + +if _RUN_LIVE_E2E: + try: + from mobilerun_core import Mobilerun + except ModuleNotFoundError as exc: + raise pytest.UsageError( + "The opted-in Grok Android E2E suite requires the public " + "mobilerun-core package in this Python environment. Install " + f"`{sys.executable} -m pip install \"mobilerun-core[local]\"` " + "and retry." + ) from exc +else: + Mobilerun = None + +pytestmark = pytest.mark.skipif( + not _RUN_LIVE_E2E, + reason="set MOBILERUN_RUN_ANDROID_E2E=1 to run live Grok Android E2E tests", +) + +_TIMEOUT_TASK = ( + "Open Android Settings, find Screen timeout, report the currently visible " + "timeout value, then return to the Home screen. Do not change any setting." +) +_ANDROID_VERSION_TASK = ( + "Open Android Settings, go to About phone, and report the visible Android " + "version, then return to the Home screen. Do not change any setting." +) + + +@dataclass(frozen=True) +class _Scenario: + scenario_id: str + provider: str + auth_mode: str + task: str + mode_flags: tuple[str, ...] + expected_state: str + + +@dataclass +class _OwnedEmulator: + binary: Path + artifact_root: Path + process: subprocess.Popen[str] | None = None + log_handle: TextIO | None = None + launch_mode: str = "" + + +@dataclass(frozen=True) +class _LiveContext: + serial: str + repo_root: Path + artifact_root: Path + configured_secrets: tuple[str, ...] + emulator: _OwnedEmulator + + +@dataclass(frozen=True) +class _ScenarioDeviceState: + home_package: str + android_version: str + timeout_labels: tuple[str, ...] + pre_evidence: _CoreEvidence + + +@dataclass(frozen=True) +class _CoreEvidence: + current_app_id: str + ui_path: Path + screenshot_path: Path + + +@dataclass(frozen=True) +class _DevicePreflight: + android_version: str + sdk_level: str + portal_version_after_setup: str + portal_version_after_doctor: str + + +@dataclass(frozen=True) +class _FileFingerprint: + exists: bool + sha256: str | None + mtime_ns: int | None + + +@dataclass(frozen=True) +class _ScenarioResult: + output: str + artifact_dir: Path + trajectory_dir: Path + events: tuple[dict[str, object], ...] + device_state: _ScenarioDeviceState + post_evidence: _CoreEvidence + + +class _AllBlackCoreScreenshot(AssertionError): + """The public Mobilerun observation returned a valid but black PNG.""" + + +_API_DIRECT = _Scenario( + scenario_id="01-xai-api-direct-ui-tree", + provider="XAI", + auth_mode="api_key", + task=_TIMEOUT_TASK, + mode_flags=("--no-reasoning", "--no-vision", "--no-vision-only"), + expected_state="screen_timeout", +) +_API_REASONING_VISION = _Scenario( + scenario_id="02-xai-api-reasoning-vision-only", + provider="XAI", + auth_mode="api_key", + task=_ANDROID_VERSION_TASK, + mode_flags=("--reasoning", "--vision", "--vision-only"), + expected_state="android_version", +) +_OAUTH_DIRECT = _Scenario( + scenario_id="03-grok-oauth-direct-ui-tree", + provider="grok_oauth", + auth_mode="oauth", + task=_TIMEOUT_TASK, + mode_flags=("--no-reasoning", "--no-vision", "--no-vision-only"), + expected_state="screen_timeout", +) +_OAUTH_REASONING_VISION = _Scenario( + scenario_id="04-grok-oauth-reasoning-vision-only", + provider="grok_oauth", + auth_mode="oauth", + task=_ANDROID_VERSION_TASK, + mode_flags=("--reasoning", "--vision", "--vision-only"), + expected_state="android_version", +) + + +def _adb(serial: str, *args: str, timeout: float = 30.0) -> str: + completed = subprocess.run( + ["adb", "-s", serial, *args], + check=False, + capture_output=True, + text=True, + timeout=timeout, + ) + if completed.returncode != 0: + pytest.fail( + f"adb {' '.join(args)} failed with exit code {completed.returncode}: " + f"{completed.stderr.strip()}" + ) + return completed.stdout.strip() + + +def _adb_readiness_probe(serial: str, *args: str) -> tuple[int, str]: + """Best-effort ADB probe while the restored emulator is still stabilizing.""" + + try: + completed = subprocess.run( + ["adb", "-s", serial, *args], + check=False, + capture_output=True, + text=True, + timeout=10, + ) + except subprocess.TimeoutExpired: + return -1, "" + return completed.returncode, completed.stdout.strip() + + +def _port_is_open(port: int) -> bool: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe: + probe.settimeout(0.25) + return probe.connect_ex(("127.0.0.1", port)) == 0 + + +def _resolve_emulator_binary() -> Path: + candidates: list[Path] = [] + for env_name in ("ANDROID_SDK_ROOT", "ANDROID_HOME"): + sdk_root = os.environ.get(env_name) + if sdk_root: + candidates.append(Path(sdk_root) / "emulator" / "emulator") + candidates.append( + Path.home() / "Library" / "Android" / "sdk" / "emulator" / "emulator" + ) + discovered = shutil.which("emulator") + if discovered: + candidates.append(Path(discovered)) + for candidate in candidates: + if candidate.is_file() and os.access(candidate, os.X_OK): + return candidate.resolve() + pytest.fail( + "Android emulator binary not found; set ANDROID_SDK_ROOT or ANDROID_HOME" + ) + + +def _emulator_command(binary: Path, *, load_snapshot: bool) -> list[str]: + command = [ + str(binary), + "-avd", + _AVD_NAME, + "-port", + str(_EMULATOR_PORT), + ] + if load_snapshot: + command.extend(("-snapshot", _SNAPSHOT_NAME)) + else: + # The recovery path intentionally preserves AVD data. + command.append("-no-snapshot-load") + command.extend(("-no-snapshot-save", "-no-boot-anim", "-no-audio")) + assert "-wipe-data" not in command + return command + + +def _wait_for_boot( + serial: str, + artifact_dir: Path, + *, + process: subprocess.Popen[str] | None, + boot_label: str, +) -> None: + """Wait for the owned emulator and Android package-manager readiness.""" + + started_at = time.monotonic() + last_boot_completed = "" + last_boot_animation = "" + package_manager_ready = False + deadline = started_at + _BOOT_TIMEOUT_SECONDS + while time.monotonic() < deadline: + if process is not None and process.poll() is not None: + raise RuntimeError( + f"owned emulator exited during {boot_label} boot " + f"with status {process.returncode}" + ) + boot_returncode, last_boot_completed = _adb_readiness_probe( + serial, "shell", "getprop", "sys.boot_completed" + ) + _, last_boot_animation = _adb_readiness_probe( + serial, "shell", "getprop", "init.svc.bootanim" + ) + package_returncode, package_output = _adb_readiness_probe( + serial, + "shell", + "cmd", + "package", + "list", + "packages", + "android", + ) + package_manager_ready = ( + package_returncode == 0 and "package:android" in package_output + ) + if ( + boot_returncode == 0 + and last_boot_completed == "1" + # With the required ``-no-boot-anim`` launch flag, Android 16 may + # leave this service property empty instead of reporting stopped. + and last_boot_animation in {"", "stopped"} + and package_manager_ready + ): + (artifact_dir / f"boot-readiness-{boot_label}.json").write_text( + json.dumps( + { + "avd": _AVD_NAME, + "serial": serial, + "port": _EMULATOR_PORT, + "boot_label": boot_label, + "sys.boot_completed": last_boot_completed, + "init.svc.bootanim": last_boot_animation, + "package_manager_ready": package_manager_ready, + "elapsed_seconds": round(time.monotonic() - started_at, 3), + }, + indent=2, + ), + encoding="utf-8", + ) + return + time.sleep(2.0) + + raise RuntimeError( + "the restored emulator did not finish booting " + f"(sys.boot_completed={last_boot_completed!r}, " + f"bootanim={last_boot_animation!r}, " + f"package_manager_ready={package_manager_ready})" + ) + + +def _owned_avd_name(serial: str) -> str: + return _adb(serial, "emu", "avd", "name").splitlines()[0].strip() + + +def _launch_owned_emulator( + emulator: _OwnedEmulator, + *, + load_snapshot: bool, + evidence_dir: Path, +) -> None: + mode = "snapshot" if load_snapshot else "cold-boot" + command = _emulator_command(emulator.binary, load_snapshot=load_snapshot) + (evidence_dir / f"emulator-launch-{mode}.json").write_text( + json.dumps( + { + "command": command, + "avd": _AVD_NAME, + "serial": _EXPECTED_SERIAL, + "port": _EMULATOR_PORT, + "snapshot": _SNAPSHOT_NAME if load_snapshot else None, + "wipe_data": False, + }, + indent=2, + ), + encoding="utf-8", + ) + emulator.log_handle = (emulator.artifact_root / "emulator.log").open( + "a", encoding="utf-8" + ) + emulator.process = subprocess.Popen( + command, + stdin=subprocess.DEVNULL, + stdout=emulator.log_handle, + stderr=subprocess.STDOUT, + text=True, + start_new_session=True, + ) + emulator.launch_mode = mode + + +def _stop_owned_emulator(emulator: _OwnedEmulator) -> None: + """Stop only the process this session launched on emulator-5558.""" + + process = emulator.process + if process is not None and process.poll() is None: + avd_probe = _adb_readiness_probe( + _EXPECTED_SERIAL, "emu", "avd", "name" + )[1] + if avd_probe.splitlines()[:1] == [_AVD_NAME]: + try: + subprocess.run( + ["adb", "-s", _EXPECTED_SERIAL, "emu", "kill"], + check=False, + capture_output=True, + text=True, + timeout=20, + ) + except subprocess.TimeoutExpired: + pass + try: + process.wait(timeout=30) + except subprocess.TimeoutExpired: + process.terminate() + try: + process.wait(timeout=10) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=10) + if emulator.log_handle is not None: + emulator.log_handle.close() + emulator.process = None + emulator.log_handle = None + + +def _wait_for_port_release(timeout: float = 30.0) -> None: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if not _port_is_open(_EMULATOR_PORT): + return + time.sleep(0.5) + raise RuntimeError(f"TCP port {_EMULATOR_PORT} did not become free") + + +def _cold_boot_owned_emulator( + emulator: _OwnedEmulator, evidence_dir: Path +) -> None: + _stop_owned_emulator(emulator) + _wait_for_port_release() + _launch_owned_emulator( + emulator, load_snapshot=False, evidence_dir=evidence_dir + ) + _wait_for_boot( + _EXPECTED_SERIAL, + evidence_dir, + process=emulator.process, + boot_label="cold-boot", + ) + if _owned_avd_name(_EXPECTED_SERIAL) != _AVD_NAME: + raise RuntimeError("cold-booted emulator reported the wrong AVD name") + + +def _start_session_emulator(artifact_root: Path) -> _OwnedEmulator: + if _port_is_open(_EMULATOR_PORT): + pytest.fail( + f"TCP port {_EMULATOR_PORT} is already occupied; the E2E suite " + "will not reuse or terminate an emulator it does not own" + ) + binary = _resolve_emulator_binary() + listed = subprocess.run( + [str(binary), "-list-avds"], + check=False, + capture_output=True, + text=True, + timeout=30, + ) + if listed.returncode != 0 or _AVD_NAME not in listed.stdout.splitlines(): + pytest.fail(f"required Android AVD {_AVD_NAME!r} is not installed") + + evidence_dir = artifact_root / "emulator-session" + evidence_dir.mkdir() + emulator = _OwnedEmulator(binary=binary, artifact_root=artifact_root) + _launch_owned_emulator( + emulator, load_snapshot=True, evidence_dir=evidence_dir + ) + try: + _wait_for_boot( + _EXPECTED_SERIAL, + evidence_dir, + process=emulator.process, + boot_label="snapshot-launch", + ) + if _owned_avd_name(_EXPECTED_SERIAL) != _AVD_NAME: + raise RuntimeError("snapshot-launched emulator reported the wrong AVD name") + except RuntimeError as snapshot_error: + (evidence_dir / "snapshot-launch-fallback.json").write_text( + json.dumps( + { + "fallback": "cold-boot", + "reason": str(snapshot_error), + "wipe_data": False, + }, + indent=2, + ), + encoding="utf-8", + ) + try: + _cold_boot_owned_emulator(emulator, evidence_dir) + except RuntimeError as cold_error: + _stop_owned_emulator(emulator) + pytest.fail(f"could not boot required AVD {_AVD_NAME!r}: {cold_error}") + return emulator + + +def _restore_clean_snapshot( + emulator: _OwnedEmulator, serial: str, artifact_dir: Path +) -> str: + """Reset a scenario, cold-booting without wipe-data if snapshot load fails.""" + + try: + completed = subprocess.run( + [ + "adb", + "-s", + serial, + "emu", + "avd", + "snapshot", + "load", + _SNAPSHOT_NAME, + ], + check=False, + capture_output=True, + text=True, + timeout=120, + ) + except subprocess.TimeoutExpired: + completed = None + + response = ( + (completed.stdout + "\n" + completed.stderr).strip() + if completed is not None + else "snapshot load timed out" + ) + (artifact_dir / "snapshot-restore.txt").write_text( + response + "\n", encoding="utf-8" + ) + snapshot_confirmed = ( + completed is not None + and completed.returncode == 0 + and not re.search(r"(?im)^KO\b", response) + and bool(re.search(r"(?im)^OK\b", response)) + ) + if snapshot_confirmed: + try: + _wait_for_boot( + serial, + artifact_dir, + process=emulator.process, + boot_label="snapshot-restore", + ) + return "snapshot" + except RuntimeError as snapshot_error: + response = f"{response}\nreadiness failure: {snapshot_error}" + + (artifact_dir / "snapshot-load-fallback.json").write_text( + json.dumps( + { + "fallback": "cold-boot", + "reason": response, + "wipe_data": False, + }, + indent=2, + ), + encoding="utf-8", + ) + try: + _cold_boot_owned_emulator(emulator, artifact_dir) + except RuntimeError as cold_error: + pytest.fail( + "snapshot load and no-wipe-data cold-boot fallback both failed: " + f"{cold_error}" + ) + return "cold_boot_fallback" + + +def _focused_package(serial: str) -> str: + window_dump = _adb(serial, "shell", "dumpsys", "window", "windows") + activity_dump = _adb(serial, "shell", "dumpsys", "activity", "activities") + ui_dump = _adb(serial, "exec-out", "uiautomator", "dump", "/dev/tty") + for output, patterns in ( + ( + window_dump, + ( + r"mCurrentFocus=Window\{[^\n]*\s([A-Za-z0-9_.]+)/", + r"mFocusedApp=.*\s([A-Za-z0-9_.]+)/", + ), + ), + ( + activity_dump, + ( + r"mResumedActivity:[^\n]*\s([A-Za-z0-9_.]+)/", + r"topResumedActivity=[^\n]*\s([A-Za-z0-9_.]+)/", + ), + ), + (ui_dump, (r'package="([A-Za-z0-9_.]+)"',)), + ): + for pattern in patterns: + match = re.search(pattern, output) + if match: + return match.group(1) + pytest.fail("could not determine the emulator's focused Android package") + + +def _timeout_labels(milliseconds: str) -> tuple[str, ...]: + try: + value = int(milliseconds.strip()) + except ValueError: + pytest.fail("screen_off_timeout was not an integer") + + if value <= 0 or value >= 2_000_000_000: + return ("never",) + + seconds = value // 1000 + if seconds == 1: + return ("1 second", "1 sec") + if seconds < 60: + return (f"{seconds} seconds", f"{seconds} sec") + + minutes = seconds // 60 + if seconds % 60 == 0: + if minutes == 1: + return ("1 minute", "1 min") + return (f"{minutes} minutes", f"{minutes} min") + return (f"{seconds} seconds", f"{seconds} sec") + + +def _credential_values(value: object, *, key: str = "") -> set[str]: + """Extract only secret-shaped values; never include metadata in failures.""" + + found: set[str] = set() + if isinstance(value, dict): + for nested_key, nested_value in value.items(): + found.update(_credential_values(nested_value, key=str(nested_key))) + elif isinstance(value, list): + for nested_value in value: + found.update(_credential_values(nested_value, key=key)) + elif isinstance(value, str) and len(value) >= 8: + normalized_key = key.casefold().replace("_", "").replace("-", "") + if "token" in normalized_key or normalized_key in { + "apikey", + "secret", + "password", + }: + found.add(value) + return found + + +def _load_grok_oauth_secrets() -> set[str]: + from mobilerun.config_manager.credential_paths import GROK_OAUTH_CREDENTIAL_PATH + + credential_path = Path(GROK_OAUTH_CREDENTIAL_PATH) + if not credential_path.is_file(): + pytest.fail( + "Grok OAuth credentials are required; run `mobilerun configure grok` " + "before the opted-in E2E suite" + ) + try: + profiles = json.loads(credential_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + pytest.fail("MobileRun's OAuth credential file is missing or malformed") + + grok_profile = profiles.get("grokOauth") if isinstance(profiles, dict) else None + if not isinstance(grok_profile, dict) or not ( + grok_profile.get("accessToken") or grok_profile.get("refreshToken") + ): + pytest.fail( + "the grokOauth slot is missing; run `mobilerun configure grok` " + "before the opted-in E2E suite" + ) + return _credential_values(grok_profile) + + +def _file_fingerprint(path: Path) -> _FileFingerprint: + """Capture only existence, SHA-256, and mtime; never parse credential data.""" + + try: + before = path.stat() + except FileNotFoundError: + return _FileFingerprint(exists=False, sha256=None, mtime_ns=None) + except OSError: + pytest.fail("could not stat the external Grok CLI credential file") + + digest = hashlib.sha256() + try: + with path.open("rb") as credential_file: + for chunk in iter(lambda: credential_file.read(1024 * 1024), b""): + digest.update(chunk) + after = path.stat() + except OSError: + pytest.fail("could not hash the external Grok CLI credential file") + if before.st_mtime_ns != after.st_mtime_ns: + pytest.fail("the external Grok CLI credential changed while being hashed") + return _FileFingerprint( + exists=True, + sha256=digest.hexdigest(), + mtime_ns=after.st_mtime_ns, + ) + + +def _write_e2e_config( + path: Path, serial: str, trajectory_path: Path +) -> None: + import yaml + + from mobilerun.config_manager import MobileConfig + from mobilerun.config_manager.migrations import CURRENT_VERSION + + config = MobileConfig() + config.agent.max_steps = 30 + config.agent.streaming = False + config.agent.app_cards.enabled = False + config.device.serial = serial + config.device.platform = "android" + config.device.use_tcp = True + config.device.portal_mode = "required" + config.device.auto_setup = False + config.telemetry.enabled = False + config.tracing.enabled = False + config.logging.debug = False + config.logging.rich_text = False + config.logging.save_trajectory = "action" + config.logging.trajectory_gifs = False + config.logging.trajectory_path = str(trajectory_path) + config.mcp.enabled = False + + payload = config.to_dict() + payload["_version"] = CURRENT_VERSION + path.write_text( + yaml.safe_dump(payload, default_flow_style=False, sort_keys=False), + encoding="utf-8", + ) + + +def _portal_version(serial: str) -> str: + output = _adb( + serial, + "shell", + "content", + "query", + "--uri", + "content://com.mobilerun.portal/version", + ) + match = re.search(r"\bresult=(\{.*\})\s*$", output) + if not match: + pytest.fail("Portal content provider did not return version evidence") + try: + payload = json.loads(match.group(1)) + except json.JSONDecodeError: + pytest.fail("Portal content provider returned malformed version evidence") + if payload.get("status") != "success": + pytest.fail("Portal content provider version query was not successful") + version = payload.get("result") or payload.get("data") + if not isinstance(version, str) or not version.strip(): + pytest.fail("Portal content provider returned an empty version") + return version.strip() + + +def _run_public_cli( + repo_root: Path, + serial: str, + *, + args: tuple[str, ...], + secrets: tuple[str, ...], + output_path: Path, + timeout: float, +) -> str: + environment = os.environ.copy() + environment.update( + { + "ANDROID_SERIAL": serial, + "MOBILERUN_TELEMETRY_ENABLED": "false", + "DROIDRUN_TELEMETRY_ENABLED": "false", + } + ) + try: + completed = subprocess.run( + [sys.executable, "-m", "mobilerun", *args], + cwd=repo_root, + env=environment, + check=False, + capture_output=True, + text=True, + timeout=timeout, + ) + except subprocess.TimeoutExpired: + pytest.fail(f"`mobilerun {args[0]}` exceeded its E2E timeout") + output = _ANSI_ESCAPE.sub("", completed.stdout + "\n" + completed.stderr) + _assert_credentials_redacted(output, secrets) + output_path.write_text(output, encoding="utf-8") + assert completed.returncode == 0, _sanitized_tail(output, secrets) + return output + + +def _setup_portal( + repo_root: Path, + serial: str, + *, + secrets: tuple[str, ...], + artifact_dir: Path, +) -> None: + last_output = "" + for attempt in range(1, 4): + output_name = ( + "mobilerun-setup.txt" + if attempt == 1 + else f"mobilerun-setup-retry-{attempt - 1}.txt" + ) + last_output = _run_public_cli( + repo_root, + serial, + args=("setup", "--device", serial), + secrets=secrets, + output_path=artifact_dir / output_name, + timeout=5 * 60, + ) + assert "setup complete!" in last_output.casefold(), ( + "`mobilerun setup` did not report successful Portal setup: " + + _sanitized_tail(last_output, secrets) + ) + assert "setup failed" not in last_output.casefold(), ( + "`mobilerun setup` reported failure: " + + _sanitized_tail(last_output, secrets) + ) + if ( + "did not become responsive" not in last_output.casefold() + and _portal_accessibility_enabled(serial) + ): + return + if attempt < 3: + time.sleep(5.0) + pytest.fail( + "Portal did not become responsive with accessibility enabled after " + "three public `mobilerun setup` attempts: " + + _sanitized_tail(last_output, secrets) + ) + + +def _ping_portal( + repo_root: Path, + serial: str, + *, + use_tcp: bool, + secrets: tuple[str, ...], + artifact_dir: Path, +) -> None: + mode_flag = "--tcp" if use_tcp else "--no-tcp" + mode_name = "tcp" if use_tcp else "content" + output = _run_public_cli( + repo_root, + serial, + args=("ping", "--device", serial, mode_flag, "--no-debug"), + secrets=secrets, + output_path=artifact_dir / f"mobilerun-ping-{mode_name}.txt", + timeout=60, + ) + mode = "TCP" if use_tcp else "content-provider" + assert "portal is installed and accessible" in output.casefold(), ( + f"{mode} Portal ping did not report success: " + + _sanitized_tail(output, secrets) + ) + + +def _doctor_portal( + repo_root: Path, + serial: str, + *, + secrets: tuple[str, ...], + artifact_dir: Path, +) -> None: + output = _run_public_cli( + repo_root, + serial, + args=("doctor", "--device", serial, "--no-debug"), + secrets=secrets, + output_path=artifact_dir / "mobilerun-doctor.txt", + timeout=5 * 60, + ) + assert "mobilerun doctor" in output.casefold(), ( + "`mobilerun doctor` did not start correctly" + ) + assert re.search(r"(?mi)^\s*Portal Version\s{2,}.*$", output), ( + "`mobilerun doctor` did not report its Portal Version check" + ) + required_checks = ( + "Device", + "Portal", + "Accessibility", + "Content Provider", + "State (content)", + "Screenshot (content)", + "TCP Mode", + "State (tcp)", + "Screenshot (tcp)", + ) + # Rich wraps long doctor rows according to the terminal width (for + # example, State (content) can continue on the next line). Bound each + # section by the following row so a checkmark from a later row cannot make + # an earlier failed row pass accidentally. + row_boundaries = (*required_checks, "Keyboard") + for index, check_name in enumerate(required_checks): + next_name = row_boundaries[index + 1] + row = re.search( + rf"(?ms)^\s*{re.escape(check_name)}\s{{2,}}.*?" + rf"(?=^\s*{re.escape(next_name)}\s{{2,}})", + output, + ) + assert row and "โœ“" in row.group(0), ( + f"`mobilerun doctor` did not pass its {check_name} check" + ) + assert not re.search(r"(?mi)^\s*\d+ issue\(s\):", output), ( + "`mobilerun doctor` reported one or more failing checks: " + + _sanitized_tail(output, secrets) + ) + + +@pytest.fixture(scope="session") +def live_context() -> _LiveContext: + serial = os.environ.get("ANDROID_SERIAL") + if serial != _EXPECTED_SERIAL: + pytest.fail( + "the opted-in Grok E2E suite requires " + f"ANDROID_SERIAL={_EXPECTED_SERIAL} exactly" + ) + if shutil.which("adb") is None: + pytest.fail("adb is required for the opted-in Android E2E suite") + + api_key = os.environ.get("XAI_API_KEY", "") + if not api_key: + pytest.fail("XAI_API_KEY is required for the opted-in Grok API E2E cases") + configured_secrets = _load_grok_oauth_secrets() + configured_secrets.add(api_key) + secrets = tuple(configured_secrets) + + repo_root = Path(__file__).resolve().parents[2] + configured_artifact_root = os.environ.get("MOBILERUN_GROK_E2E_ARTIFACTS") + if configured_artifact_root: + artifact_base = Path(configured_artifact_root).expanduser().resolve() + artifact_base.mkdir(parents=True, exist_ok=True) + artifact_root = artifact_base / ( + f"run-{time.strftime('%Y%m%d-%H%M%S')}-{os.getpid()}-{time.time_ns()}" + ) + artifact_root.mkdir() + else: + artifact_root = Path( + tempfile.mkdtemp(prefix="mobilerun-grok-android16-") + ).resolve() + + external_grok_auth = Path.home() / ".grok" / "auth.json" + external_grok_before = _file_fingerprint(external_grok_auth) + emulator = _start_session_emulator(artifact_root) + context = _LiveContext( + serial=serial, + repo_root=repo_root, + artifact_root=artifact_root, + configured_secrets=secrets, + emulator=emulator, + ) + try: + yield context + finally: + _stop_owned_emulator(emulator) + external_grok_after = _file_fingerprint(external_grok_auth) + hash_unchanged = external_grok_before.sha256 == external_grok_after.sha256 + mtime_unchanged = ( + external_grok_before.mtime_ns == external_grok_after.mtime_ns + ) + existence_unchanged = ( + external_grok_before.exists == external_grok_after.exists + ) + (artifact_root / "external-grok-cli-auth-invariant.json").write_text( + json.dumps( + { + "path": "~/.grok/auth.json", + "existence_unchanged": existence_unchanged, + "sha256_unchanged": hash_unchanged, + "mtime_unchanged": mtime_unchanged, + "contents_parsed": False, + }, + indent=2, + ), + encoding="utf-8", + ) + if not (existence_unchanged and hash_unchanged and mtime_unchanged): + pytest.fail( + "MobileRun's live Grok suite modified the separate Grok CLI " + "credential at ~/.grok/auth.json" + ) + + +_ANSI_ESCAPE = re.compile(r"\x1b\[[0-?]*[ -/]*[@-~]") +_TOKEN_ASSIGNMENT = re.compile( + r"(?i)(?:access[_-]?token|refresh[_-]?token|authorization)\s*[:=]\s*" + r"[\"']?(?:bearer\s+)?[A-Za-z0-9._~+/-]{12,}" +) + + +def _assert_credentials_redacted(output: str, secrets: tuple[str, ...]) -> None: + for secret in secrets: + assert secret not in output, "Mobilerun output exposed a configured secret" + assert not _TOKEN_ASSIGNMENT.search(output), ( + "Mobilerun output exposed a token-shaped credential assignment" + ) + + +def _assert_artifacts_redacted(root: Path, secrets: tuple[str, ...]) -> None: + """Scan every artifact as bytes and every UTF-8 artifact as text.""" + + for artifact in sorted(path for path in root.rglob("*") if path.is_file()): + data = artifact.read_bytes() + relative_path = artifact.relative_to(root) + for secret in secrets: + assert secret.encode("utf-8") not in data, ( + f"artifact {relative_path} exposed a configured secret" + ) + try: + text = data.decode("utf-8") + except UnicodeDecodeError: + continue + assert not _TOKEN_ASSIGNMENT.search(text), ( + f"text artifact {relative_path} exposed a token-shaped credential" + ) + + +def _sanitized_tail(output: str, secrets: tuple[str, ...]) -> str: + sanitized = output + for secret in secrets: + sanitized = sanitized.replace(secret, "[REDACTED]") + sanitized = _TOKEN_ASSIGNMENT.sub("credential=[REDACTED]", sanitized) + return sanitized[-4000:] + + +def _jsonable(value: Any) -> Any: + if hasattr(value, "model_dump"): + return value.model_dump(mode="json") + if isinstance(value, dict): + return {str(key): _jsonable(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_jsonable(item) for item in value] + if isinstance(value, (str, int, float, bool)) or value is None: + return value + return str(value) + + +def _connect_core_device(context: _LiveContext) -> Any: + if Mobilerun is None: + pytest.fail( + "mobilerun_core.Mobilerun is unavailable in the live E2E environment" + ) + try: + device = Mobilerun().connect( + context.serial, + backend="local-android-adb", + portal_mode="required", + ) + except Exception as exc: + pytest.fail(f"public mobilerun-core device connection failed: {exc}") + assert device.capabilities.get("platform") == "android" + return device + + +def _capture_core_evidence( + context: _LiveContext, + artifact_dir: Path, + *, + label: str, + ensure_home: bool, +) -> _CoreEvidence: + """Use only the public mobile-harness surface for control and evidence.""" + + device = _connect_core_device(context) + if ensure_home: + # Snapshot restores may preserve an asleep display. Use only the + # public mobile-harness key/swipe surface for the ordinary wake, + # unlock, and Home sequence before observing the framebuffer. + device.key("wakeup") + time.sleep(0.5) + width, height = device.screen_size() + device.swipe(width // 2, height * 4 // 5, width // 2, height // 5, ms=350) + time.sleep(0.5) + device.key("home") + device.wait_for_idle(timeout=5.0) + current_app_id = device.current_app_id() or "" + ui = device.ui() + screenshot = base64.b64decode(device.screenshot(hide_overlay=True)) + ui_path = artifact_dir / f"{label}-core-ui.json" + screenshot_path = artifact_dir / f"{label}-core.png" + ui_path.write_text( + json.dumps(_jsonable(ui), ensure_ascii=False, indent=2), + encoding="utf-8", + ) + screenshot_path.write_bytes(screenshot) + assert screenshot.startswith(b"\x89PNG\r\n\x1a\n"), ( + "mobilerun-core evidence screenshot was not a PNG" + ) + from PIL import Image + + with Image.open(io.BytesIO(screenshot)) as image: + if image.convert("RGB").getbbox() is None: + raise _AllBlackCoreScreenshot( + "public mobilerun-core returned an all-black screenshot after " + "wake, unlock, and Home" + ) + return _CoreEvidence( + current_app_id=current_app_id, + ui_path=ui_path, + screenshot_path=screenshot_path, + ) + + +def _portal_accessibility_enabled(serial: str) -> bool: + enabled_services = _adb( + serial, + "shell", + "settings", + "get", + "secure", + "enabled_accessibility_services", + ) + return "com.mobilerun.portal" in enabled_services + + +def _setup_and_check_device( + context: _LiveContext, + artifact_dir: Path, +) -> _DevicePreflight: + """Run public setup/ping/doctor and independent readiness assertions.""" + + _setup_portal( + context.repo_root, + context.serial, + secrets=context.configured_secrets, + artifact_dir=artifact_dir, + ) + + android_version = _adb( + context.serial, "shell", "getprop", "ro.build.version.release" + ) + assert android_version == "16", ( + f"{context.serial} must run Android 16 during scenario preflight " + f"(reported {android_version!r})" + ) + sdk_level = _adb( + context.serial, "shell", "getprop", "ro.build.version.sdk" + ) + assert sdk_level == "36", ( + f"{context.serial} must use Android SDK 36 during scenario preflight " + f"(reported {sdk_level!r})" + ) + + portal_package = _adb( + context.serial, "shell", "pm", "path", "com.mobilerun.portal" + ) + assert portal_package.startswith("package:"), ( + "`mobilerun setup` did not install Mobilerun Portal" + ) + portal_version_after_setup = _portal_version(context.serial) + assert re.fullmatch( + r"v?\d+(?:\.\d+){1,3}(?:[-+][A-Za-z0-9.-]+)?", + portal_version_after_setup, + ), "Portal content provider returned an invalid version string" + assert _portal_accessibility_enabled(context.serial), ( + "`mobilerun setup` did not enable Portal accessibility" + ) + + _ping_portal( + context.repo_root, + context.serial, + use_tcp=False, + secrets=context.configured_secrets, + artifact_dir=artifact_dir, + ) + _ping_portal( + context.repo_root, + context.serial, + use_tcp=True, + secrets=context.configured_secrets, + artifact_dir=artifact_dir, + ) + _doctor_portal( + context.repo_root, + context.serial, + secrets=context.configured_secrets, + artifact_dir=artifact_dir, + ) + + portal_version_after_doctor = _portal_version(context.serial) + assert re.fullmatch( + r"v?\d+(?:\.\d+){1,3}(?:[-+][A-Za-z0-9.-]+)?", + portal_version_after_doctor, + ), "Portal content provider returned an invalid post-doctor version string" + assert _portal_accessibility_enabled(context.serial), ( + "Portal accessibility was not enabled after `mobilerun doctor`" + ) + return _DevicePreflight( + android_version=android_version, + sdk_level=sdk_level, + portal_version_after_setup=portal_version_after_setup, + portal_version_after_doctor=portal_version_after_doctor, + ) + + +def _prepare_scenario_device( + context: _LiveContext, + scenario: _Scenario, + artifact_dir: Path, +) -> _ScenarioDeviceState: + """Restore, boot, set up, and fully validate the device for one scenario.""" + + initial_reset_mode = _restore_clean_snapshot( + context.emulator, context.serial, artifact_dir + ) + final_reset_mode = initial_reset_mode + graphics_fallback_reason: str | None = None + preflight_artifact_dir = artifact_dir + preflight = _setup_and_check_device(context, preflight_artifact_dir) + try: + pre_evidence = _capture_core_evidence( + context, + preflight_artifact_dir, + label="pre-task", + ensure_home=True, + ) + except _AllBlackCoreScreenshot as exc: + if initial_reset_mode != "snapshot": + raise + graphics_fallback_reason = str(exc) + fallback_dir = artifact_dir / "snapshot-graphics-cold-boot" + fallback_dir.mkdir() + (artifact_dir / "snapshot-graphics-fallback.json").write_text( + json.dumps( + { + "snapshot_attempted": True, + "snapshot_restored": True, + "failure": "all_black_public_core_framebuffer", + "reason": graphics_fallback_reason, + "fallback": "cold-boot", + "wipe_data": False, + "rerun_after_fallback": [ + "mobilerun setup", + "mobilerun ping --no-tcp", + "mobilerun ping --tcp", + "mobilerun doctor", + "public mobilerun-core evidence", + ], + }, + indent=2, + ), + encoding="utf-8", + ) + try: + _cold_boot_owned_emulator(context.emulator, fallback_dir) + except RuntimeError as cold_error: + pytest.fail( + "snapshot framebuffer was all-black and the no-wipe-data " + f"cold-boot fallback failed: {cold_error}" + ) + preflight_artifact_dir = fallback_dir + preflight = _setup_and_check_device(context, preflight_artifact_dir) + try: + pre_evidence = _capture_core_evidence( + context, + preflight_artifact_dir, + label="pre-task", + ensure_home=True, + ) + except _AllBlackCoreScreenshot: + pytest.fail( + "public mobilerun-core framebuffer remained all-black after " + "the no-wipe-data cold-boot fallback" + ) + final_reset_mode = "snapshot_graphics_cold_boot_fallback" + + home_package = pre_evidence.current_app_id + assert home_package, "mobilerun-core could not identify the Home package" + timeout_labels = _timeout_labels( + _adb( + context.serial, + "shell", + "settings", + "get", + "system", + "screen_off_timeout", + ) + ) + (artifact_dir / "preflight.json").write_text( + json.dumps( + { + "scenario_id": scenario.scenario_id, + "serial": context.serial, + "avd": _AVD_NAME, + "emulator_port": _EMULATOR_PORT, + "snapshot": _SNAPSHOT_NAME, + "snapshot_attempted": True, + "snapshot_restored": initial_reset_mode == "snapshot", + "initial_reset_mode": initial_reset_mode, + "reset_mode": final_reset_mode, + "snapshot_graphics_fallback_reason": graphics_fallback_reason, + "cold_boot_without_wipe_data": final_reset_mode + in { + "cold_boot_fallback", + "snapshot_graphics_cold_boot_fallback", + }, + "boot_completed": True, + "mobilerun_setup": "passed", + "android_version": preflight.android_version, + "sdk_level": preflight.sdk_level, + "portal_version_after_setup": preflight.portal_version_after_setup, + "portal_version_after_doctor": preflight.portal_version_after_doctor, + "portal_accessibility": "enabled", + "content_provider_ping": "passed", + "tcp_ping": "passed", + "mobilerun_doctor": "passed", + "portal_mode": "required", + "auto_setup_during_task": False, + }, + indent=2, + ), + encoding="utf-8", + ) + _assert_artifacts_redacted(artifact_dir, context.configured_secrets) + return _ScenarioDeviceState( + home_package=home_package, + android_version=preflight.android_version, + timeout_labels=timeout_labels, + pre_evidence=pre_evidence, + ) + + +def _load_and_assert_trajectory(trajectory_root: Path) -> tuple[Path, tuple[dict[str, object], ...]]: + trajectory_files = sorted(trajectory_root.glob("*/trajectory.json")) + assert len(trajectory_files) == 1, ( + "each Grok scenario must produce exactly one trajectory.json artifact" + ) + trajectory_file = trajectory_files[0] + try: + raw_events = json.loads(trajectory_file.read_text(encoding="utf-8")) + except json.JSONDecodeError: + pytest.fail("trajectory.json was not valid JSON") + assert isinstance(raw_events, list) and raw_events, ( + "trajectory.json must contain recorded MobileRun events" + ) + events = tuple(event for event in raw_events if isinstance(event, dict)) + assert len(events) == len(raw_events), "trajectory events must be JSON objects" + + response_events = [ + event + for event in events + if str(event.get("type", "")).endswith("ResponseEvent") + ] + assert any( + isinstance(event.get("usage"), dict) + and int(event["usage"].get("requests", 0) or 0) > 0 + and int(event["usage"].get("total_tokens", 0) or 0) > 0 + for event in response_events + ), "trajectory must include positive request and token usage on a ResponseEvent" + + assert any( + event.get("type") == "ToolExecutionEvent" and event.get("success") is True + for event in events + ), "trajectory must include a successful ToolExecutionEvent" + direct_completion = any( + event.get("type") == "ToolExecutionEvent" + and event.get("tool_name") == "complete" + and event.get("success") is True + and isinstance(event.get("tool_args"), dict) + and event["tool_args"].get("success") is True + for event in events + ) + reasoning_completion = any( + event.get("type") == "ManagerResponseEvent" + and isinstance(event.get("response"), str) + and re.search( + r"]*\bsuccess=[\"']true[\"']", + event["response"], + flags=re.IGNORECASE, + ) + for event in events + ) + assert direct_completion or reasoning_completion, ( + "trajectory must include a successful MobileRun completion result" + ) + + trajectory_dir = trajectory_file.parent + ui_state_files = sorted((trajectory_dir / "ui_states").glob("*.json")) + screenshot_files = sorted((trajectory_dir / "screenshots").glob("*.png")) + assert ui_state_files, "trajectory must include recorded UI-state artifacts" + assert screenshot_files, "trajectory must include screenshot artifacts" + for ui_state_file in ui_state_files: + try: + ui_state = json.loads(ui_state_file.read_text(encoding="utf-8")) + except json.JSONDecodeError: + pytest.fail("a recorded UI-state artifact was not valid JSON") + assert isinstance(ui_state, list), "recorded UI state must be a JSON list" + for screenshot_file in screenshot_files: + screenshot = screenshot_file.read_bytes() + assert screenshot.startswith(b"\x89PNG\r\n\x1a\n"), ( + "recorded screenshot must be a non-empty PNG artifact" + ) + return trajectory_dir, events + + +def _recorded_screen_timeout_label(trajectory_dir: Path) -> str: + """Read the visible Settings summary from independently recorded Portal UI.""" + + for ui_state_file in sorted((trajectory_dir / "ui_states").glob("*.json")): + raw_nodes = json.loads(ui_state_file.read_text(encoding="utf-8")) + if not isinstance(raw_nodes, list): + continue + for index, node in enumerate(raw_nodes): + if not isinstance(node, dict) or str(node.get("text", "")).casefold() != "screen timeout": + continue + for summary in raw_nodes[index + 1 : index + 6]: + if not isinstance(summary, dict): + continue + resource_id = str(summary.get("resourceId", "")) + label = str(summary.get("text", "")).strip() + if resource_id.endswith("id/summary") and label: + return label + pytest.fail("recorded Portal UI did not contain the visible Screen timeout summary") + + +def _run_scenario( + context: _LiveContext, scenario: _Scenario +) -> _ScenarioResult: + artifact_dir = context.artifact_root / scenario.scenario_id + artifact_dir.mkdir(parents=False, exist_ok=False) + trajectory_root = artifact_dir / "trajectories" + config_path = artifact_dir / "config.yaml" + (artifact_dir / "scenario.json").write_text( + json.dumps( + { + "scenario_id": scenario.scenario_id, + "snapshot": _SNAPSHOT_NAME, + "provider": scenario.provider, + "auth_mode": scenario.auth_mode, + "model": "grok-4.5", + "mode_flags": list(scenario.mode_flags), + "portal_mode": "required", + "auto_setup": False, + "tcp": True, + "tracing": False, + "telemetry": False, + "trajectory": "action", + }, + indent=2, + ), + encoding="utf-8", + ) + device_state = _prepare_scenario_device(context, scenario, artifact_dir) + _write_e2e_config(config_path, context.serial, trajectory_root) + command = [ + sys.executable, + "-m", + "mobilerun", + "run", + scenario.task, + "--config", + str(config_path), + "--device", + context.serial, + "--tcp", + "--provider", + scenario.provider, + "--model", + "grok-4.5", + "--steps", + "30", + "--no-stream", + "--no-tracing", + "--no-debug", + "--save-trajectory", + "action", + *scenario.mode_flags, + ] + environment = os.environ.copy() + environment.update( + { + "ANDROID_SERIAL": context.serial, + "MOBILERUN_CONFIG": str(config_path), + "MOBILERUN_TELEMETRY_ENABLED": "false", + "DROIDRUN_TELEMETRY_ENABLED": "false", + } + ) + # The natural-language agent CLI is the Grok integration under test. All + # harness-owned device actions and evidence use mobilerun_core.Mobilerun. + completed = subprocess.run( + command, + cwd=context.repo_root, + env=environment, + check=False, + capture_output=True, + text=True, + timeout=20 * 60, + ) + output = _ANSI_ESCAPE.sub("", completed.stdout + "\n" + completed.stderr) + # OAuth may rotate tokens during inference. Scan for both the credentials + # that existed before the suite and the newly persisted credential values. + runtime_secrets = set(context.configured_secrets) + runtime_secrets.update(_load_grok_oauth_secrets()) + runtime_secrets.add(os.environ["XAI_API_KEY"]) + secrets = tuple(runtime_secrets) + _assert_credentials_redacted(output, secrets) + (artifact_dir / "mobilerun-run.txt").write_text(output, encoding="utf-8") + post_evidence = _capture_core_evidence( + context, + artifact_dir, + label="post-task", + ensure_home=False, + ) + _assert_artifacts_redacted(artifact_dir, secrets) + assert completed.returncode == 0, _sanitized_tail( + output, secrets + ) + trajectory_dir, events = _load_and_assert_trajectory(trajectory_root) + return _ScenarioResult( + output=output, + artifact_dir=artifact_dir, + trajectory_dir=trajectory_dir, + events=events, + device_state=device_state, + post_evidence=post_evidence, + ) + + +def _assert_scenario_state( + context: _LiveContext, + scenario: _Scenario, + result: _ScenarioResult, +) -> None: + normalized_output = " ".join(result.output.casefold().split()) + assert result.post_evidence.current_app_id == result.device_state.home_package, ( + "Grok did not finish the task on the Home screen" + ) + assert _focused_package(context.serial) == result.device_state.home_package, ( + "independent foreground diagnostics did not confirm the Home screen" + ) + if scenario.expected_state == "screen_timeout": + visible_label = _recorded_screen_timeout_label(result.trajectory_dir) + assert " ".join(visible_label.casefold().split()) in normalized_output, ( + "Grok did not report the emulator's visible Screen timeout value" + ) + return + + assert result.device_state.android_version.casefold() in normalized_output, ( + "Grok did not report the emulator's Android version" + ) + assert ( + _adb(context.serial, "shell", "getprop", "ro.build.version.release") + == result.device_state.android_version + ) + + +def test_xai_api_direct_ui_tree_screen_timeout(live_context: _LiveContext) -> None: + result = _run_scenario(live_context, _API_DIRECT) + _assert_scenario_state(live_context, _API_DIRECT, result) + + +def test_xai_api_reasoning_vision_only_android_version( + live_context: _LiveContext, +) -> None: + result = _run_scenario(live_context, _API_REASONING_VISION) + _assert_scenario_state(live_context, _API_REASONING_VISION, result) + + +def test_grok_oauth_direct_ui_tree_screen_timeout( + live_context: _LiveContext, +) -> None: + result = _run_scenario(live_context, _OAUTH_DIRECT) + _assert_scenario_state(live_context, _OAUTH_DIRECT, result) + + +def test_grok_oauth_reasoning_vision_only_android_version( + live_context: _LiveContext, +) -> None: + result = _run_scenario(live_context, _OAUTH_REASONING_VISION) + _assert_scenario_state(live_context, _OAUTH_REASONING_VISION, result) diff --git a/tests/test_grok_android_e2e_helpers.py b/tests/test_grok_android_e2e_helpers.py new file mode 100644 index 00000000..d6926c2a --- /dev/null +++ b/tests/test_grok_android_e2e_helpers.py @@ -0,0 +1,174 @@ +from __future__ import annotations + +import base64 +import io +import json +from pathlib import Path +from typing import Any + +import pytest +from PIL import Image + +from tests.e2e import test_grok_android16 as android_e2e + + +class _BlackScreenshotDevice: + capabilities = {"platform": "android"} + + def __init__(self, screenshot: bytes) -> None: + self._screenshot = screenshot + self.actions: list[tuple[Any, ...]] = [] + + def key(self, name: str) -> None: + self.actions.append(("key", name)) + + def screen_size(self) -> tuple[int, int]: + return 1080, 2400 + + def swipe(self, *args: Any, **kwargs: Any) -> None: + self.actions.append(("swipe", *args, kwargs)) + + def wait_for_idle(self, *, timeout: float) -> bool: + self.actions.append(("wait_for_idle", timeout)) + return True + + def current_app_id(self) -> str: + return "com.android.launcher3" + + def ui(self) -> dict[str, object]: + return {"phone_state": {"package_name": "com.android.launcher3"}} + + def screenshot(self, *, hide_overlay: bool) -> str: + assert hide_overlay is True + return base64.b64encode(self._screenshot).decode("ascii") + + +def _black_png() -> bytes: + output = io.BytesIO() + Image.new("RGB", (4, 4), color="black").save(output, format="PNG") + return output.getvalue() + + +def _context(tmp_path: Path) -> android_e2e._LiveContext: + emulator = android_e2e._OwnedEmulator( + binary=tmp_path / "emulator", + artifact_root=tmp_path, + ) + return android_e2e._LiveContext( + serial="emulator-5558", + repo_root=tmp_path, + artifact_root=tmp_path, + configured_secrets=(), + emulator=emulator, + ) + + +def test_public_core_black_png_raises_graphics_readiness_error( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + device = _BlackScreenshotDevice(_black_png()) + monkeypatch.setattr(android_e2e, "_connect_core_device", lambda context: device) + monkeypatch.setattr(android_e2e.time, "sleep", lambda seconds: None) + + with pytest.raises( + android_e2e._AllBlackCoreScreenshot, + match="all-black screenshot", + ): + android_e2e._capture_core_evidence( + _context(tmp_path), + tmp_path, + label="pre-task", + ensure_home=True, + ) + + assert (tmp_path / "pre-task-core.png").read_bytes() == device._screenshot + assert device.actions[0] == ("key", "wakeup") + assert ("key", "home") in device.actions + + +def test_snapshot_black_framebuffer_uses_no_wipe_cold_boot_and_rechecks( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + artifact_dir = tmp_path / "scenario" + artifact_dir.mkdir() + context = _context(tmp_path) + setup_dirs: list[Path] = [] + capture_dirs: list[Path] = [] + cold_boot_dirs: list[Path] = [] + preflight = android_e2e._DevicePreflight( + android_version="16", + sdk_level="36", + portal_version_after_setup="0.7.1", + portal_version_after_doctor="0.7.1", + ) + recovered = android_e2e._CoreEvidence( + current_app_id="com.android.launcher3", + ui_path=artifact_dir / "recovered-ui.json", + screenshot_path=artifact_dir / "recovered.png", + ) + + monkeypatch.setattr( + android_e2e, + "_restore_clean_snapshot", + lambda emulator, serial, artifacts: "snapshot", + ) + + def setup_and_check( + unused_context: android_e2e._LiveContext, + artifacts: Path, + ) -> android_e2e._DevicePreflight: + setup_dirs.append(artifacts) + return preflight + + monkeypatch.setattr(android_e2e, "_setup_and_check_device", setup_and_check) + + def capture( + unused_context: android_e2e._LiveContext, + artifacts: Path, + *, + label: str, + ensure_home: bool, + ) -> android_e2e._CoreEvidence: + assert label == "pre-task" + assert ensure_home is True + capture_dirs.append(artifacts) + if len(capture_dirs) == 1: + raise android_e2e._AllBlackCoreScreenshot("all-black public frame") + return recovered + + monkeypatch.setattr(android_e2e, "_capture_core_evidence", capture) + monkeypatch.setattr( + android_e2e, + "_cold_boot_owned_emulator", + lambda emulator, artifacts: cold_boot_dirs.append(artifacts), + ) + monkeypatch.setattr(android_e2e, "_adb", lambda serial, *args: "30000") + monkeypatch.setattr( + android_e2e, + "_timeout_labels", + lambda timeout_ms: ("30 seconds",), + ) + + result = android_e2e._prepare_scenario_device( + context, + android_e2e._API_DIRECT, + artifact_dir, + ) + + fallback_dir = artifact_dir / "snapshot-graphics-cold-boot" + assert setup_dirs == [artifact_dir, fallback_dir] + assert capture_dirs == [artifact_dir, fallback_dir] + assert cold_boot_dirs == [fallback_dir] + assert result.pre_evidence is recovered + fallback = json.loads( + (artifact_dir / "snapshot-graphics-fallback.json").read_text() + ) + assert fallback["failure"] == "all_black_public_core_framebuffer" + assert fallback["wipe_data"] is False + summary = json.loads((artifact_dir / "preflight.json").read_text()) + assert summary["snapshot_attempted"] is True + assert summary["snapshot_restored"] is True + assert summary["reset_mode"] == "snapshot_graphics_cold_boot_fallback" + assert summary["cold_boot_without_wipe_data"] is True diff --git a/tests/test_grok_api.py b/tests/test_grok_api.py new file mode 100644 index 00000000..6b009513 --- /dev/null +++ b/tests/test_grok_api.py @@ -0,0 +1,553 @@ +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("grok", "api_key") + + assert variant.id == "XAI" + assert variant.runtime_provider_name == "XAI" + assert variant.default_model == "grok-4.5" + assert list_models_for_variant("grok", "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("grok", "oauth") + + assert variant.id == "grok_oauth" + assert variant.runtime_provider_name == "grok_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("grok", auth_mode, model_alias) == "grok-4.5" + + +@pytest.mark.parametrize("alias", ("grok", "xai", "x.ai", "XAI")) +def test_grok_runtime_aliases_select_xai(alias: str) -> None: + assert normalize_provider_name(alias) == "XAI" + + +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("grok", "api_key") + profile = create_profile_for_variant( + variant, + SetupSelection( + family_id="grok", + 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 == "grok" + 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("grok", "api_key") + profile = create_profile_for_variant( + variant, + SetupSelection( + family_id="grok", + 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("grok", model="grok-4.5") + + assert llm.api_key == "xai-runtime-key" + + +@pytest.mark.parametrize("alias", ("grok", "xai", "x.ai", "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_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..aed93df8 --- /dev/null +++ b/tests/test_grok_cli.py @@ -0,0 +1,178 @@ +import json +from types import SimpleNamespace + +import pytest +from click.testing import CliRunner + +import mobilerun.cli.configure_wizard as configure_wizard +import mobilerun.cli.main as cli_main +from mobilerun.cli.configure_wizard import ConfigureWizardCallbacks +from mobilerun.config_manager import MobileConfig + + +def test_grok_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), "grok_oauth" + ) + + credential_path.write_text( + json.dumps({"openaiOauth": {"access": "unrelated"}}), + encoding="utf-8", + ) + assert not configure_wizard._oauth_credentials_present( + str(credential_path), "grok_oauth" + ) + + +def test_wizard_prepares_grok_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="grok_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_grok_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", + "grok", + "--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_help_advertises_grok_provider_and_login_command() -> None: + runner = CliRunner() + + configure_help = runner.invoke(cli_main.cli, ["configure", "--help"]) + grok_help = runner.invoke(cli_main.cli, ["configure", "grok", "--help"]) + + assert configure_help.exit_code == 0 + assert "grok" in configure_help.output.lower() + assert grok_help.exit_code == 0 + assert "--device-code" in grok_help.output + assert "native xAI OAuth" in grok_help.output + + +@pytest.mark.parametrize( + ("auth_mode", "expected_provider"), + (("api_key", "XAI"), ("oauth", "grok_oauth")), +) +def test_exact_grok_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 Grok 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", "grok", "--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, "grok", auth_mode, "grok-4.5")} diff --git a/tests/test_grok_oauth.py b/tests/test_grok_oauth.py new file mode 100644 index 00000000..526394ef --- /dev/null +++ b/tests/test_grok_oauth.py @@ -0,0 +1,1496 @@ +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.config_manager import 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"} + + +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_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"}, + ) + jwks_client = SimpleNamespace( + get_signing_key_from_jwt=lambda _: SimpleNamespace(key=public_key) + ) + validator = GrokIDTokenValidator(jwks_client) + + assert validator.validate(token, nonce="expected")["sub"] == "user" + with pytest.raises(jwt.InvalidTokenError, match="nonce"): + validator.validate(token, nonce="wrong") + + +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 +): + 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) + 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) + + 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] + + +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] = [] + + def handler(request: httpx.Request) -> httpx.Response: + nonlocal attempts + attempts += 1 + 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] + 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", + ) + + assert attempts == 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, tmp_path: Path): + requests: list[httpx.Request] = [] + + 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": "https://accounts.x.ai/device?code=ABCD-1234", + "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) + 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, 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 + + +def test_device_token_poll_retries_connect_and_5xx_failures(tmp_path: Path): + poll_attempts = 0 + delays: list[float] = [] + + def handler(request: httpx.Request) -> httpx.Response: + nonlocal poll_attempts + if request.url.path.endswith("/device/code"): + 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: + raise httpx.ConnectError("temporary connect failure", request=request) + if poll_attempts == 2: + return httpx.Response(503, text="sensitive-upstream-body") + return httpx.Response( + 200, + json={ + "access_token": "access", + "refresh_token": "refresh", + "expires_in": 3600, + }, + ) + + 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] + retry_backoff_seconds=(0.01, 0.02), + sleep=delays.append, + ) + + credentials = GrokOAuth(oauth_session_manager=manager).login( + device_code=True, + timeout_seconds=10, + ) + + assert poll_attempts == 3 + assert delays == [0.01, 0.02] + assert credentials.access_token == "access" + + +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 Grok 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 Grok 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_llm_picker.py b/tests/test_llm_picker.py index 93ef3c64..0cc91415 100644 --- a/tests/test_llm_picker.py +++ b/tests/test_llm_picker.py @@ -125,6 +125,7 @@ async def parse_async(**kwargs: Any) -> Any: for payload in (sync_payload, async_payload): assert {"temperature", "top_p"}.isdisjoint(payload) assert payload["max_output_tokens"] == 32 + assert payload["tool_choice"] == "none" @pytest.mark.parametrize( diff --git a/tests/test_usage.py b/tests/test_usage.py index 6e7e6bcd..df43aaff 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("GrokOAuth", 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" }, From 52dcd52d8a68e269e139b5ea475fac48089aabdf Mon Sep 17 00:00:00 2001 From: "rasul.osmanbayli" Date: Thu, 13 Aug 2026 11:49:39 +0400 Subject: [PATCH 2/4] Tighten XAI provider integration --- README.md | 17 +- docs/guides/cli.mdx | 39 +- docs/quickstart.mdx | 9 +- docs/sdk/configuration.mdx | 20 +- mobilerun/agent/providers/grok.py | 4 +- mobilerun/agent/providers/registry.py | 10 +- mobilerun/agent/providers/setup_service.py | 3 +- mobilerun/agent/usage.py | 2 +- mobilerun/agent/utils/llm_picker.py | 15 +- mobilerun/agent/utils/oauth/grok_oauth_llm.py | 44 +- mobilerun/cli/configure_wizard.py | 21 +- mobilerun/cli/main.py | 20 +- mobilerun/cli/oauth_actions.py | 2 +- mobilerun/config_example.yaml | 18 - .../config_manager/auth_profile_store.py | 6 +- tests/e2e/test_grok_android16.py | 1554 ----------------- tests/test_grok_android_e2e_helpers.py | 174 -- tests/test_grok_api.py | 49 +- tests/test_grok_cli.py | 137 +- tests/test_grok_oauth.py | 16 +- tests/test_llm_picker.py | 1 - tests/test_usage.py | 2 +- 22 files changed, 248 insertions(+), 1915 deletions(-) delete mode 100644 tests/e2e/test_grok_android16.py delete mode 100644 tests/test_grok_android_e2e_helpers.py diff --git a/README.md b/README.md index 12b3ee23..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, Grok, 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 @@ -115,21 +115,6 @@ 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`, `XAI_API_KEY`, or `MINIMAX_API_KEY`. -Grok supports either an xAI API key or a native xAI OAuth login managed by -Mobilerun: - -```bash -# Direct xAI API -export XAI_API_KEY=your-key -mobilerun configure --provider grok --auth-mode api_key --model grok-4.5 - -# Native Grok OAuth (use --device-code on SSH/headless machines) -mobilerun configure grok -``` - -Mobilerun stores this login in its own `grokOauth` credential slot. It does not -invoke the Grok CLI or read `~/.grok/auth.json`. - ### 4. Run your first command ```bash diff --git a/docs/guides/cli.mdx b/docs/guides/cli.mdx index b0dff4d8..27f96e08 100644 --- a/docs/guides/cli.mdx +++ b/docs/guides/cli.mdx @@ -161,7 +161,6 @@ 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` | -| grok_oauth | Included by default | Mobilerun-managed xAI OAuth credentials | | OpenAILike | Included by default | Varies by provider | | OpenRouter | Included by default | `OPENROUTER_API_KEY` | | Ollama | Included by default | None (local) | @@ -169,44 +168,26 @@ mobilerun run "Enable 2FA" \ | DeepSeek | Included by default | `DEEPSEEK_API_KEY` | | MiniMax | Included by default | `MINIMAX_API_KEY` | -### Grok API key and OAuth +### xAI API key and OAuth -Grok 4.5 is available through the public xAI API and through a Grok account. -For public API billing, set `XAI_API_KEY` or paste the key in the configure -wizard: +Configure XAI with an API key or OAuth: ```bash +# API key mobilerun configure \ - --provider grok \ + --provider XAI \ --auth-mode api_key \ --model grok-4.5 -``` -For Grok OAuth, let Mobilerun run its native xAI login flow and save the result -in the shared `auth-profiles.json` store: +# OAuth through provider options +mobilerun configure --provider XAI --auth-mode oauth --model grok-4.5 -```bash -# Browser login -mobilerun configure grok +# OAuth shortcut +mobilerun configure xai # Device-code login for SSH/headless hosts -mobilerun configure grok --device-code --no-browser -``` - -You can also choose **Grok** and **oauth** in `mobilerun configure`. The -resulting profile uses provider `grok_oauth`; API-key profiles use provider -`XAI`. OAuth credentials are stored under `grokOauth` with owner-only file -permissions. This credential is separate from `~/.grok/auth.json`: Mobilerun -does not invoke the Grok CLI, read its credential file, or treat an OAuth -failure as permission to fall back to `XAI_API_KEY`. - - - Grok OAuth currently interoperates with xAI's public Grok Build OAuth client - and subscription proxy contract. xAI does not document that contract as a - stable third-party integration, so xAI may require a fresh login or a - compatibility update in a future release. Use `XAI` with an API key when you - need the documented public API boundary. - +mobilerun configure xai --device-code --no-browser +``` ### MiniMax endpoints and credentials diff --git a/docs/quickstart.mdx b/docs/quickstart.mdx index 3a9117f0..5cabe91b 100644 --- a/docs/quickstart.mdx +++ b/docs/quickstart.mdx @@ -105,15 +105,10 @@ export OPENAI_API_KEY=your-api-key-here # For Anthropic Claude export ANTHROPIC_API_KEY=your-api-key-here -# For xAI Grok +# For xAI export XAI_API_KEY=your-api-key-here ``` -To use your Grok account instead of an API key, run -`mobilerun configure grok` (add `--device-code --no-browser` on a headless -host). Mobilerun runs its own xAI OAuth flow and never invokes the Grok CLI or -reads `~/.grok/auth.json`. - ### Run Your First Command via CLI Now you're ready to control your device with natural language: @@ -133,7 +128,7 @@ mobilerun run "Find a contact named John and send him an email" --reasoning ``` **Common CLI flags:** -- `--provider` - LLM provider (GoogleGenAI, OpenAI, XAI, grok_oauth, 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 8774cb47..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, XAI, grok_oauth, 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 @@ -641,24 +641,6 @@ export MINIMAX_API_KEY=your-key export MOBILERUN_CONFIG=/path/to/config.yaml # Custom config path ``` -Grok 4.5 supports both the public xAI API and Grok OAuth. Use the `XAI` -runtime provider for API-key calls: - -```bash -mobilerun configure \ - --provider grok \ - --auth-mode api_key \ - --model grok-4.5 -``` - -For OAuth, run `mobilerun configure grok` (or add -`--device-code --no-browser` on a headless host). Mobilerun performs its own -xAI OAuth flow, stores the credential under `grokOauth` in the shared platform -`credentials/auth-profiles.json` file, and generates a `grok_oauth` profile. -The Grok CLI is not required or invoked, and its `~/.grok/auth.json` remains -separate. API-key profiles use `apiKeys.xai` / `XAI_API_KEY` and the `XAI` -runtime provider. - MiniMax uses `https://api.minimax.io/v1` for global accounts and `https://api.minimaxi.com/v1` for Mainland China accounts. The configure wizard asks for the matching region and selects `MiniMax-M3` by default: diff --git a/mobilerun/agent/providers/grok.py b/mobilerun/agent/providers/grok.py index e350a98f..d6c77bc9 100644 --- a/mobilerun/agent/providers/grok.py +++ b/mobilerun/agent/providers/grok.py @@ -24,7 +24,7 @@ def normalize_grok_model_id(model: object) -> str: - """Normalize public xAI/Grok aliases to MobileRun's canonical model id.""" + """Normalize public xAI/Grok aliases to Mobilerun's canonical model id.""" model_id = str(model or "").strip() if model_id.startswith("xai/"): @@ -54,7 +54,7 @@ def sanitize_grok_responses_kwargs( # 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. + # 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): diff --git a/mobilerun/agent/providers/registry.py b/mobilerun/agent/providers/registry.py index f9f74cdd..7d9d69f4 100644 --- a/mobilerun/agent/providers/registry.py +++ b/mobilerun/agent/providers/registry.py @@ -142,8 +142,8 @@ ), ), ProviderFamilySpec( - id="grok", - display_name="Grok", + id="xai", + display_name="XAI", variants=( ProviderVariantSpec( id="XAI", @@ -155,8 +155,8 @@ base_url=XAI_API_BASE, ), ProviderVariantSpec( - id="grok_oauth", - runtime_provider_name="grok_oauth", + id="xai_oauth", + runtime_provider_name="xai_oauth", auth_mode="oauth", default_model=GROK_DEFAULT_MODEL, models=GROK_MODELS, @@ -316,7 +316,7 @@ def normalize_model_id_for_variant( if family_id == "openai": candidate = OPENAI_MODEL_ALIASES.get(candidate, candidate) - elif family_id == "grok": + elif family_id == "xai": candidate = normalize_grok_model_id(candidate) if candidate in allowed_model_ids: diff --git a/mobilerun/agent/providers/setup_service.py b/mobilerun/agent/providers/setup_service.py index 02989635..25a3a1de 100644 --- a/mobilerun/agent/providers/setup_service.py +++ b/mobilerun/agent/providers/setup_service.py @@ -151,6 +151,7 @@ def create_profile_for_variant( resolved_model, ) kwargs: dict[str, str | int] = dict(DEFAULT_KWARGS_BY_VARIANT.get(variant.id, {})) + env_slot = VARIANT_ENV_KEY_SLOT.get(variant.id) runtime_provider_name = ( variant.runtime_transport_provider_name or variant.runtime_provider_name ) @@ -181,7 +182,7 @@ def create_profile_for_variant( else None ), credential_path=selection.credential_path or variant.credential_path, - kwargs=kwargs, + 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 1cdb3d9d..36a5c39d 100644 --- a/mobilerun/agent/usage.py +++ b/mobilerun/agent/usage.py @@ -29,7 +29,7 @@ "MobilerunOpenAIResponses": "OpenAIResponses", "openai_responses_llm": "OpenAIResponses", "GrokOAuth": "OpenAIResponses", - "grok_oauth": "OpenAIResponses", + "xai_oauth": "OpenAIResponses", "Ollama_llm": "Ollama", } diff --git a/mobilerun/agent/utils/llm_picker.py b/mobilerun/agent/utils/llm_picker.py index 45bc1039..94fc4d4c 100644 --- a/mobilerun/agent/utils/llm_picker.py +++ b/mobilerun/agent/utils/llm_picker.py @@ -57,9 +57,7 @@ "openai_like": "OpenAILike", "zai": "ZAI", "z.ai": "ZAI", - "grok": "XAI", "xai": "XAI", - "x.ai": "XAI", } ZAI_GLOBAL_API_BASE = "https://api.z.ai/api/paas/v4" @@ -529,8 +527,8 @@ def load_llm(provider_name: str, model: str | None = None, **kwargs: Any) -> LLM model = normalize_model_id_for_variant("openai", "api_key", model) elif provider_name == "XAI": model = normalize_grok_model_id(model) - elif provider_name == "grok_oauth": - model = normalize_model_id_for_variant("grok", "oauth", 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 @@ -560,7 +558,7 @@ 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 == "grok_oauth": + 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}) @@ -596,7 +594,7 @@ def load_llm(provider_name: str, model: str | None = None, **kwargs: Any) -> LLM if provider_name == "XAI": import os - # MobileRun's reasoning mode selects its agent architecture. It does + # 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") @@ -609,9 +607,8 @@ def load_llm(provider_name: str, model: str | None = None, **kwargs: Any) -> LLM ) kwargs["api_key"] = api_key - # The public runtime aliases (`grok`, `xai`, and `x.ai`) should be - # useful without a separately generated profile as well. Keep their - # implicit model aligned with the first-class provider catalog. + # 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 diff --git a/mobilerun/agent/utils/oauth/grok_oauth_llm.py b/mobilerun/agent/utils/oauth/grok_oauth_llm.py index 218f4dc9..c7cbc2d7 100644 --- a/mobilerun/agent/utils/oauth/grok_oauth_llm.py +++ b/mobilerun/agent/utils/oauth/grok_oauth_llm.py @@ -1,6 +1,6 @@ """xAI subscription OAuth transport for the Responses API. -Credentials are owned by MobileRun and stored in its shared auth profile. The +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. """ @@ -173,16 +173,16 @@ class GrokOAuthCredentials: @classmethod def from_payload(cls, payload: dict[str, Any]) -> "GrokOAuthCredentials": if payload.get("type") != "oauth" or payload.get("provider") != "xai-grok": - raise ValueError("Grok OAuth profile has an unexpected credential type.") + 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("Grok OAuth profile is missing accessToken.") + 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("Grok OAuth profile has an unexpected issuer.") + raise ValueError("XAI OAuth profile has an unexpected issuer.") if client_id != DEFAULT_GROK_OAUTH_CLIENT_ID: - raise ValueError("Grok OAuth profile has an unexpected clientId.") + raise ValueError("XAI OAuth profile has an unexpected clientId.") refresh = payload.get("refreshToken") raw_expiry = payload.get("expiresAt") @@ -198,7 +198,7 @@ def from_payload(cls, payload: dict[str, Any]) -> "GrokOAuthCredentials": ) token_type = str(payload.get("tokenType") or "Bearer") if token_type.lower() != "bearer": - raise ValueError("Grok OAuth profile has an unsupported tokenType.") + 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, @@ -271,7 +271,7 @@ def validate(self, token: str, *, nonce: str | None) -> dict[str, Any]: class GrokOAuthSessionManager: - """Load, exchange, and cross-process-refresh MobileRun's xAI session.""" + """Load, exchange, and cross-process-refresh Mobilerun's xAI session.""" def __init__( self, @@ -420,7 +420,7 @@ def _post_token( error = None if refresh_request and error in {"invalid_grant", "invalid_client"}: raise GrokOAuthReloginRequired( - "Grok OAuth refresh was rejected; re-login is required." + "XAI OAuth refresh was rejected; re-login is required." ) raise GrokOAuthError( f"xAI token request failed ({error or response.status_code})." @@ -453,7 +453,7 @@ def exchange_authorization_code( def _refresh(self, credentials: GrokOAuthCredentials) -> GrokOAuthCredentials: if not credentials.refresh_token: raise ValueError( - "No Grok OAuth refresh token is available. Run MobileRun's Grok login." + "No XAI OAuth refresh token is available. Run `mobilerun configure xai`." ) payload = self._post_token( { @@ -477,7 +477,7 @@ def get_valid_credentials( rejected_access_token: str | None = None, ) -> GrokOAuthCredentials: with self._thread_lock: - # Keep the file lock across refresh so separate MobileRun processes + # 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) @@ -488,7 +488,7 @@ def get_valid_credentials( ) if credentials is None: raise ValueError( - "No Grok OAuth credentials found. Run MobileRun's Grok login." + "No XAI OAuth credentials found. Run `mobilerun configure xai`." ) another_writer_refreshed = ( @@ -566,7 +566,7 @@ async def async_auth_flow(self, request: httpx.Request): # type: ignore[no-unty class GrokOAuth(OpenAIResponses): - """LlamaIndex Responses adapter backed only by MobileRun's xAI OAuth slot.""" + """LlamaIndex Responses adapter backed only by Mobilerun's xAI OAuth slot.""" @classmethod def class_name(cls) -> str: @@ -589,7 +589,7 @@ def __init__( model = normalize_grok_model_id(model) if model not in GROK_MODELS: raise ValueError( - f"Model {model!r} is not supported with Grok OAuth. " + f"Model {model!r} is not supported with XAI OAuth. " f"Use {', '.join(GROK_MODELS)}." ) path = ( @@ -773,11 +773,11 @@ def login( device_code: bool = False, ) -> GrokOAuthCredentials: if callback_host != DEFAULT_GROK_OAUTH_CALLBACK_HOST: - raise ValueError("Grok OAuth callback_host must be 127.0.0.1.") + raise ValueError("XAI OAuth callback_host must be 127.0.0.1.") if callback_port != 0: - raise ValueError("Grok OAuth callback_port must be OS-assigned (0).") + raise ValueError("XAI OAuth callback_port must be OS-assigned (0).") if callback_path != DEFAULT_GROK_OAUTH_CALLBACK_PATH: - raise ValueError("Grok OAuth callback_path must be /callback.") + raise ValueError("XAI OAuth callback_path must be /callback.") if device_code or _is_headless_environment(): return self._login_device_code(timeout_seconds=timeout_seconds) @@ -808,9 +808,9 @@ def do_GET(self) -> None: # noqa: N802 self.send_header("Content-Type", "text/html; charset=utf-8") self.end_headers() self.wfile.write( - b"MobileRun Grok login complete. You may close this tab." + b"Mobilerun XAI login complete. You may close this tab." if ok - else b"MobileRun Grok login failed. Return to the terminal." + else b"Mobilerun XAI login failed. Return to the terminal." ) done.set() @@ -839,15 +839,15 @@ def log_message(self, format: str, *args: Any) -> None: # noqa: A003 if open_browser: webbrowser.open(authorization_url) if not done.wait(timeout=max(0.0, timeout_seconds)): - raise TimeoutError("Grok OAuth login timed out waiting for callback.") + raise TimeoutError("XAI OAuth login timed out waiting for callback.") 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("Grok OAuth callback state mismatch.") + raise GrokOAuthError("XAI OAuth callback state mismatch.") if not result["code"]: - raise GrokOAuthError("Grok OAuth callback did not contain a 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, @@ -942,7 +942,7 @@ def _login_device_code( f"xAI device token request failed ({error or token_response.status_code})." ) manager.sleep(min(interval, max(0.0, deadline - time.monotonic()))) - raise TimeoutError("Grok OAuth device authorization timed out.") + raise TimeoutError("XAI OAuth device authorization timed out.") # Descriptive alias for callers that prefer the full class name. diff --git a/mobilerun/cli/configure_wizard.py b/mobilerun/cli/configure_wizard.py index f11c0dca..8afe3da2 100644 --- a/mobilerun/cli/configure_wizard.py +++ b/mobilerun/cli/configure_wizard.py @@ -108,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", @@ -243,7 +244,7 @@ def _prompt_base_url_for_variant(variant: Any) -> str: "gemini_oauth_code_assist": "geminiAntigravityOauth", "openai_oauth": "openaiOauth", "anthropic_oauth": "claudeAiOauth", - "grok_oauth": "grokOauth", + "xai_oauth": "grokOauth", } @@ -323,9 +324,9 @@ def _prepare_variant_auth( callbacks.run_gemini_oauth_login( credential_path=credential_path, model=selected_model ) - elif variant.id == "grok_oauth" and credential_path: + elif variant.id == "xai_oauth" and credential_path: if callbacks.run_grok_oauth_login is None: - raise RuntimeError("Grok OAuth login callback is not configured.") + raise RuntimeError("XAI OAuth login callback is not configured.") callbacks.run_grok_oauth_login( credential_path=credential_path, model=selected_model, @@ -556,6 +557,10 @@ 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 @@ -691,11 +696,11 @@ def run_configure_wizard( if model_is_fixed: state.selected_model = model - # A fixed provider should enter that provider's flow immediately. Missing - # auth/model/key values remain interactive, but the supplied provider (and - # auth mode, when present) must not be discarded by the top-level 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: + 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 34ecc1a1..c76d101d 100644 --- a/mobilerun/cli/main.py +++ b/mobilerun/cli/main.py @@ -26,8 +26,6 @@ ping_portal_tcp, setup_portal, ) -from mobilerun_core_local.driver.ios import discover_ios_device, validate_ios_portal_url -from mobilerun_core_local.driver.visual_remote import VISUAL_REMOTE_CONNECTION from rich.console import Console from rich.panel import Panel from rich.text import Text @@ -68,6 +66,8 @@ from mobilerun.log_handlers import CLILogHandler, configure_logging from mobilerun.macro.cli import macro_cli from mobilerun.telemetry import print_telemetry_message +from mobilerun_core_local.driver.ios import discover_ios_device, validate_ios_portal_url +from mobilerun_core_local.driver.visual_remote import VISUAL_REMOTE_CONNECTION # Suppress all warnings warnings.filterwarnings("ignore") @@ -432,7 +432,7 @@ def _run_grok_oauth_login( model=model, **kwargs, ) - _print_oauth_login_success("Grok", credential_path) + _print_oauth_login_success("XAI", credential_path) try: @@ -460,7 +460,7 @@ def _run_grok_oauth_login( @click.option( "--provider", "-p", - help="LLM provider (OpenAI, openai_oauth, XAI, grok_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( @@ -1025,7 +1025,7 @@ async def doctor(device: str | None, debug: bool | None): "--provider", type=str, default=None, - help="Provider family (gemini, openai, anthropic, grok, ollama, openai_like, minimax, zai).", + help="Provider family (gemini, openai, anthropic, XAI, ollama, openai_like, minimax, zai).", ) @click.option( "--auth-mode", @@ -1222,12 +1222,12 @@ def configure_gemini( ) -@configure.command("grok") +@configure.command("xai") @click.option( "--credential-path", default=str(GROK_OAUTH_CREDENTIAL_PATH), show_default=True, - help="Where to store Mobilerun's Grok OAuth credentials.", + help="Where to store XAI OAuth credentials.", ) @click.option( "--model", default=None, help="Optional model override for later API calls." @@ -1249,16 +1249,16 @@ def configure_gemini( "--device-code", is_flag=True, default=False, - help="Use Grok's device-code flow for SSH or other headless environments.", + help="Use xAI's device-code flow for SSH or other headless environments.", ) -def configure_grok( +def configure_xai( credential_path: str, model: str | None, timeout: float, open_browser: bool, device_code: bool, ): - """Log in to Grok with Mobilerun's native xAI OAuth flow.""" + """Log in to XAI with Mobilerun OAuth.""" _run_grok_oauth_login( credential_path=credential_path, model=model, diff --git a/mobilerun/cli/oauth_actions.py b/mobilerun/cli/oauth_actions.py index 1c43821c..1ce171c4 100644 --- a/mobilerun/cli/oauth_actions.py +++ b/mobilerun/cli/oauth_actions.py @@ -88,7 +88,7 @@ def run_grok_oauth_login( device_code: bool = False, no_browser: bool = False, ) -> None: - """Authenticate directly with xAI and save MobileRun-owned credentials.""" + """Authenticate directly with xAI and save Mobilerun-owned credentials.""" llm = GrokOAuth( model=model or DEFAULT_GROK_MODEL, oauth_credential_path=credential_path, diff --git a/mobilerun/config_example.yaml b/mobilerun/config_example.yaml index 0f1ec3d0..0058166f 100644 --- a/mobilerun/config_example.yaml +++ b/mobilerun/config_example.yaml @@ -78,24 +78,6 @@ llm_profiles: # kwargs: # optional kwargs, add api_key in kwargs if not already in .env # max_tokens: 8192 # - # Grok API-key example (reads apiKeys.xai / XAI_API_KEY): - # provider: XAI - # provider_family: grok - # auth_mode: api_key - # model: grok-4.5 - # api_key_source: auto - # base_url: https://api.x.ai/v1 - # api_base: https://api.x.ai/v1 - # kwargs: - # context_window: 500000 - # - # Grok OAuth example (first run `mobilerun configure grok`): - # provider: grok_oauth - # provider_family: grok - # auth_mode: oauth - # model: grok-4.5 - # credential_path: /droidrun/credentials/auth-profiles.json - # # Ollama example โ€” max_tokens is translated to Ollama's num_predict, and # context_window controls num_ctx (defaults to 32768; -1 = model maximum, # which preallocates the full KV cache and can spill to CPU): diff --git a/mobilerun/config_manager/auth_profile_store.py b/mobilerun/config_manager/auth_profile_store.py index 72fc6d4f..e436091e 100644 --- a/mobilerun/config_manager/auth_profile_store.py +++ b/mobilerun/config_manager/auth_profile_store.py @@ -1,4 +1,4 @@ -"""Safe shared storage for MobileRun authentication profiles. +"""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 @@ -22,6 +22,7 @@ class AuthProfileFormatError(ValueError): _T = TypeVar("_T") +_FCHMOD: Callable[[int, int], None] | None = getattr(os, "fchmod", None) class AuthProfileTransaction(AbstractContextManager["AuthProfileTransaction"]): @@ -129,7 +130,8 @@ def _write_unlocked(self, profile: dict[str, Any]) -> None: suffix=".tmp", ) tmp_path = Path(raw_tmp_path) - os.fchmod(fd, 0o600) + 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) diff --git a/tests/e2e/test_grok_android16.py b/tests/e2e/test_grok_android16.py deleted file mode 100644 index 9fd4b5db..00000000 --- a/tests/e2e/test_grok_android16.py +++ /dev/null @@ -1,1554 +0,0 @@ -"""Opt-in Grok API/OAuth E2E coverage on the dedicated Android 16 emulator. - -These tests exercise the natural-language agent through public ``mobilerun -run`` and use the mobile harness's public ``mobilerun_core.Mobilerun`` surface -for harness-owned device actions and evidence. They are excluded from normal -test runs because they consume live xAI credentials and mutate the emulator's -foreground UI. To run them, set both of the following explicitly:: - - MOBILERUN_RUN_ANDROID_E2E=1 - ANDROID_SERIAL=emulator-5558 - -The host must provide AVD ``mobilerun_agent_bench_api36`` with snapshot -``mobilerun_eval_clean_api36``. The session claims only port 5558, launches and -owns that exact emulator, and tears down only ``emulator-5558``. Every scenario -restores the snapshot, runs ``mobilerun setup``, and verifies the Portal before -starting the task. Agent auto-setup remains disabled during the task itself. -""" - -from __future__ import annotations - -import base64 -import hashlib -import io -import json -import os -import re -import shutil -import socket -import subprocess -import sys -import tempfile -import time -from dataclasses import dataclass -from pathlib import Path -from typing import Any, TextIO - -import pytest - -_EXPECTED_SERIAL = "emulator-5558" -_EMULATOR_PORT = 5558 -_AVD_NAME = "mobilerun_agent_bench_api36" -_SNAPSHOT_NAME = "mobilerun_eval_clean_api36" -_BOOT_TIMEOUT_SECONDS = 240.0 -_RUN_LIVE_E2E = os.environ.get("MOBILERUN_RUN_ANDROID_E2E") == "1" - -if _RUN_LIVE_E2E: - try: - from mobilerun_core import Mobilerun - except ModuleNotFoundError as exc: - raise pytest.UsageError( - "The opted-in Grok Android E2E suite requires the public " - "mobilerun-core package in this Python environment. Install " - f"`{sys.executable} -m pip install \"mobilerun-core[local]\"` " - "and retry." - ) from exc -else: - Mobilerun = None - -pytestmark = pytest.mark.skipif( - not _RUN_LIVE_E2E, - reason="set MOBILERUN_RUN_ANDROID_E2E=1 to run live Grok Android E2E tests", -) - -_TIMEOUT_TASK = ( - "Open Android Settings, find Screen timeout, report the currently visible " - "timeout value, then return to the Home screen. Do not change any setting." -) -_ANDROID_VERSION_TASK = ( - "Open Android Settings, go to About phone, and report the visible Android " - "version, then return to the Home screen. Do not change any setting." -) - - -@dataclass(frozen=True) -class _Scenario: - scenario_id: str - provider: str - auth_mode: str - task: str - mode_flags: tuple[str, ...] - expected_state: str - - -@dataclass -class _OwnedEmulator: - binary: Path - artifact_root: Path - process: subprocess.Popen[str] | None = None - log_handle: TextIO | None = None - launch_mode: str = "" - - -@dataclass(frozen=True) -class _LiveContext: - serial: str - repo_root: Path - artifact_root: Path - configured_secrets: tuple[str, ...] - emulator: _OwnedEmulator - - -@dataclass(frozen=True) -class _ScenarioDeviceState: - home_package: str - android_version: str - timeout_labels: tuple[str, ...] - pre_evidence: _CoreEvidence - - -@dataclass(frozen=True) -class _CoreEvidence: - current_app_id: str - ui_path: Path - screenshot_path: Path - - -@dataclass(frozen=True) -class _DevicePreflight: - android_version: str - sdk_level: str - portal_version_after_setup: str - portal_version_after_doctor: str - - -@dataclass(frozen=True) -class _FileFingerprint: - exists: bool - sha256: str | None - mtime_ns: int | None - - -@dataclass(frozen=True) -class _ScenarioResult: - output: str - artifact_dir: Path - trajectory_dir: Path - events: tuple[dict[str, object], ...] - device_state: _ScenarioDeviceState - post_evidence: _CoreEvidence - - -class _AllBlackCoreScreenshot(AssertionError): - """The public Mobilerun observation returned a valid but black PNG.""" - - -_API_DIRECT = _Scenario( - scenario_id="01-xai-api-direct-ui-tree", - provider="XAI", - auth_mode="api_key", - task=_TIMEOUT_TASK, - mode_flags=("--no-reasoning", "--no-vision", "--no-vision-only"), - expected_state="screen_timeout", -) -_API_REASONING_VISION = _Scenario( - scenario_id="02-xai-api-reasoning-vision-only", - provider="XAI", - auth_mode="api_key", - task=_ANDROID_VERSION_TASK, - mode_flags=("--reasoning", "--vision", "--vision-only"), - expected_state="android_version", -) -_OAUTH_DIRECT = _Scenario( - scenario_id="03-grok-oauth-direct-ui-tree", - provider="grok_oauth", - auth_mode="oauth", - task=_TIMEOUT_TASK, - mode_flags=("--no-reasoning", "--no-vision", "--no-vision-only"), - expected_state="screen_timeout", -) -_OAUTH_REASONING_VISION = _Scenario( - scenario_id="04-grok-oauth-reasoning-vision-only", - provider="grok_oauth", - auth_mode="oauth", - task=_ANDROID_VERSION_TASK, - mode_flags=("--reasoning", "--vision", "--vision-only"), - expected_state="android_version", -) - - -def _adb(serial: str, *args: str, timeout: float = 30.0) -> str: - completed = subprocess.run( - ["adb", "-s", serial, *args], - check=False, - capture_output=True, - text=True, - timeout=timeout, - ) - if completed.returncode != 0: - pytest.fail( - f"adb {' '.join(args)} failed with exit code {completed.returncode}: " - f"{completed.stderr.strip()}" - ) - return completed.stdout.strip() - - -def _adb_readiness_probe(serial: str, *args: str) -> tuple[int, str]: - """Best-effort ADB probe while the restored emulator is still stabilizing.""" - - try: - completed = subprocess.run( - ["adb", "-s", serial, *args], - check=False, - capture_output=True, - text=True, - timeout=10, - ) - except subprocess.TimeoutExpired: - return -1, "" - return completed.returncode, completed.stdout.strip() - - -def _port_is_open(port: int) -> bool: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe: - probe.settimeout(0.25) - return probe.connect_ex(("127.0.0.1", port)) == 0 - - -def _resolve_emulator_binary() -> Path: - candidates: list[Path] = [] - for env_name in ("ANDROID_SDK_ROOT", "ANDROID_HOME"): - sdk_root = os.environ.get(env_name) - if sdk_root: - candidates.append(Path(sdk_root) / "emulator" / "emulator") - candidates.append( - Path.home() / "Library" / "Android" / "sdk" / "emulator" / "emulator" - ) - discovered = shutil.which("emulator") - if discovered: - candidates.append(Path(discovered)) - for candidate in candidates: - if candidate.is_file() and os.access(candidate, os.X_OK): - return candidate.resolve() - pytest.fail( - "Android emulator binary not found; set ANDROID_SDK_ROOT or ANDROID_HOME" - ) - - -def _emulator_command(binary: Path, *, load_snapshot: bool) -> list[str]: - command = [ - str(binary), - "-avd", - _AVD_NAME, - "-port", - str(_EMULATOR_PORT), - ] - if load_snapshot: - command.extend(("-snapshot", _SNAPSHOT_NAME)) - else: - # The recovery path intentionally preserves AVD data. - command.append("-no-snapshot-load") - command.extend(("-no-snapshot-save", "-no-boot-anim", "-no-audio")) - assert "-wipe-data" not in command - return command - - -def _wait_for_boot( - serial: str, - artifact_dir: Path, - *, - process: subprocess.Popen[str] | None, - boot_label: str, -) -> None: - """Wait for the owned emulator and Android package-manager readiness.""" - - started_at = time.monotonic() - last_boot_completed = "" - last_boot_animation = "" - package_manager_ready = False - deadline = started_at + _BOOT_TIMEOUT_SECONDS - while time.monotonic() < deadline: - if process is not None and process.poll() is not None: - raise RuntimeError( - f"owned emulator exited during {boot_label} boot " - f"with status {process.returncode}" - ) - boot_returncode, last_boot_completed = _adb_readiness_probe( - serial, "shell", "getprop", "sys.boot_completed" - ) - _, last_boot_animation = _adb_readiness_probe( - serial, "shell", "getprop", "init.svc.bootanim" - ) - package_returncode, package_output = _adb_readiness_probe( - serial, - "shell", - "cmd", - "package", - "list", - "packages", - "android", - ) - package_manager_ready = ( - package_returncode == 0 and "package:android" in package_output - ) - if ( - boot_returncode == 0 - and last_boot_completed == "1" - # With the required ``-no-boot-anim`` launch flag, Android 16 may - # leave this service property empty instead of reporting stopped. - and last_boot_animation in {"", "stopped"} - and package_manager_ready - ): - (artifact_dir / f"boot-readiness-{boot_label}.json").write_text( - json.dumps( - { - "avd": _AVD_NAME, - "serial": serial, - "port": _EMULATOR_PORT, - "boot_label": boot_label, - "sys.boot_completed": last_boot_completed, - "init.svc.bootanim": last_boot_animation, - "package_manager_ready": package_manager_ready, - "elapsed_seconds": round(time.monotonic() - started_at, 3), - }, - indent=2, - ), - encoding="utf-8", - ) - return - time.sleep(2.0) - - raise RuntimeError( - "the restored emulator did not finish booting " - f"(sys.boot_completed={last_boot_completed!r}, " - f"bootanim={last_boot_animation!r}, " - f"package_manager_ready={package_manager_ready})" - ) - - -def _owned_avd_name(serial: str) -> str: - return _adb(serial, "emu", "avd", "name").splitlines()[0].strip() - - -def _launch_owned_emulator( - emulator: _OwnedEmulator, - *, - load_snapshot: bool, - evidence_dir: Path, -) -> None: - mode = "snapshot" if load_snapshot else "cold-boot" - command = _emulator_command(emulator.binary, load_snapshot=load_snapshot) - (evidence_dir / f"emulator-launch-{mode}.json").write_text( - json.dumps( - { - "command": command, - "avd": _AVD_NAME, - "serial": _EXPECTED_SERIAL, - "port": _EMULATOR_PORT, - "snapshot": _SNAPSHOT_NAME if load_snapshot else None, - "wipe_data": False, - }, - indent=2, - ), - encoding="utf-8", - ) - emulator.log_handle = (emulator.artifact_root / "emulator.log").open( - "a", encoding="utf-8" - ) - emulator.process = subprocess.Popen( - command, - stdin=subprocess.DEVNULL, - stdout=emulator.log_handle, - stderr=subprocess.STDOUT, - text=True, - start_new_session=True, - ) - emulator.launch_mode = mode - - -def _stop_owned_emulator(emulator: _OwnedEmulator) -> None: - """Stop only the process this session launched on emulator-5558.""" - - process = emulator.process - if process is not None and process.poll() is None: - avd_probe = _adb_readiness_probe( - _EXPECTED_SERIAL, "emu", "avd", "name" - )[1] - if avd_probe.splitlines()[:1] == [_AVD_NAME]: - try: - subprocess.run( - ["adb", "-s", _EXPECTED_SERIAL, "emu", "kill"], - check=False, - capture_output=True, - text=True, - timeout=20, - ) - except subprocess.TimeoutExpired: - pass - try: - process.wait(timeout=30) - except subprocess.TimeoutExpired: - process.terminate() - try: - process.wait(timeout=10) - except subprocess.TimeoutExpired: - process.kill() - process.wait(timeout=10) - if emulator.log_handle is not None: - emulator.log_handle.close() - emulator.process = None - emulator.log_handle = None - - -def _wait_for_port_release(timeout: float = 30.0) -> None: - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - if not _port_is_open(_EMULATOR_PORT): - return - time.sleep(0.5) - raise RuntimeError(f"TCP port {_EMULATOR_PORT} did not become free") - - -def _cold_boot_owned_emulator( - emulator: _OwnedEmulator, evidence_dir: Path -) -> None: - _stop_owned_emulator(emulator) - _wait_for_port_release() - _launch_owned_emulator( - emulator, load_snapshot=False, evidence_dir=evidence_dir - ) - _wait_for_boot( - _EXPECTED_SERIAL, - evidence_dir, - process=emulator.process, - boot_label="cold-boot", - ) - if _owned_avd_name(_EXPECTED_SERIAL) != _AVD_NAME: - raise RuntimeError("cold-booted emulator reported the wrong AVD name") - - -def _start_session_emulator(artifact_root: Path) -> _OwnedEmulator: - if _port_is_open(_EMULATOR_PORT): - pytest.fail( - f"TCP port {_EMULATOR_PORT} is already occupied; the E2E suite " - "will not reuse or terminate an emulator it does not own" - ) - binary = _resolve_emulator_binary() - listed = subprocess.run( - [str(binary), "-list-avds"], - check=False, - capture_output=True, - text=True, - timeout=30, - ) - if listed.returncode != 0 or _AVD_NAME not in listed.stdout.splitlines(): - pytest.fail(f"required Android AVD {_AVD_NAME!r} is not installed") - - evidence_dir = artifact_root / "emulator-session" - evidence_dir.mkdir() - emulator = _OwnedEmulator(binary=binary, artifact_root=artifact_root) - _launch_owned_emulator( - emulator, load_snapshot=True, evidence_dir=evidence_dir - ) - try: - _wait_for_boot( - _EXPECTED_SERIAL, - evidence_dir, - process=emulator.process, - boot_label="snapshot-launch", - ) - if _owned_avd_name(_EXPECTED_SERIAL) != _AVD_NAME: - raise RuntimeError("snapshot-launched emulator reported the wrong AVD name") - except RuntimeError as snapshot_error: - (evidence_dir / "snapshot-launch-fallback.json").write_text( - json.dumps( - { - "fallback": "cold-boot", - "reason": str(snapshot_error), - "wipe_data": False, - }, - indent=2, - ), - encoding="utf-8", - ) - try: - _cold_boot_owned_emulator(emulator, evidence_dir) - except RuntimeError as cold_error: - _stop_owned_emulator(emulator) - pytest.fail(f"could not boot required AVD {_AVD_NAME!r}: {cold_error}") - return emulator - - -def _restore_clean_snapshot( - emulator: _OwnedEmulator, serial: str, artifact_dir: Path -) -> str: - """Reset a scenario, cold-booting without wipe-data if snapshot load fails.""" - - try: - completed = subprocess.run( - [ - "adb", - "-s", - serial, - "emu", - "avd", - "snapshot", - "load", - _SNAPSHOT_NAME, - ], - check=False, - capture_output=True, - text=True, - timeout=120, - ) - except subprocess.TimeoutExpired: - completed = None - - response = ( - (completed.stdout + "\n" + completed.stderr).strip() - if completed is not None - else "snapshot load timed out" - ) - (artifact_dir / "snapshot-restore.txt").write_text( - response + "\n", encoding="utf-8" - ) - snapshot_confirmed = ( - completed is not None - and completed.returncode == 0 - and not re.search(r"(?im)^KO\b", response) - and bool(re.search(r"(?im)^OK\b", response)) - ) - if snapshot_confirmed: - try: - _wait_for_boot( - serial, - artifact_dir, - process=emulator.process, - boot_label="snapshot-restore", - ) - return "snapshot" - except RuntimeError as snapshot_error: - response = f"{response}\nreadiness failure: {snapshot_error}" - - (artifact_dir / "snapshot-load-fallback.json").write_text( - json.dumps( - { - "fallback": "cold-boot", - "reason": response, - "wipe_data": False, - }, - indent=2, - ), - encoding="utf-8", - ) - try: - _cold_boot_owned_emulator(emulator, artifact_dir) - except RuntimeError as cold_error: - pytest.fail( - "snapshot load and no-wipe-data cold-boot fallback both failed: " - f"{cold_error}" - ) - return "cold_boot_fallback" - - -def _focused_package(serial: str) -> str: - window_dump = _adb(serial, "shell", "dumpsys", "window", "windows") - activity_dump = _adb(serial, "shell", "dumpsys", "activity", "activities") - ui_dump = _adb(serial, "exec-out", "uiautomator", "dump", "/dev/tty") - for output, patterns in ( - ( - window_dump, - ( - r"mCurrentFocus=Window\{[^\n]*\s([A-Za-z0-9_.]+)/", - r"mFocusedApp=.*\s([A-Za-z0-9_.]+)/", - ), - ), - ( - activity_dump, - ( - r"mResumedActivity:[^\n]*\s([A-Za-z0-9_.]+)/", - r"topResumedActivity=[^\n]*\s([A-Za-z0-9_.]+)/", - ), - ), - (ui_dump, (r'package="([A-Za-z0-9_.]+)"',)), - ): - for pattern in patterns: - match = re.search(pattern, output) - if match: - return match.group(1) - pytest.fail("could not determine the emulator's focused Android package") - - -def _timeout_labels(milliseconds: str) -> tuple[str, ...]: - try: - value = int(milliseconds.strip()) - except ValueError: - pytest.fail("screen_off_timeout was not an integer") - - if value <= 0 or value >= 2_000_000_000: - return ("never",) - - seconds = value // 1000 - if seconds == 1: - return ("1 second", "1 sec") - if seconds < 60: - return (f"{seconds} seconds", f"{seconds} sec") - - minutes = seconds // 60 - if seconds % 60 == 0: - if minutes == 1: - return ("1 minute", "1 min") - return (f"{minutes} minutes", f"{minutes} min") - return (f"{seconds} seconds", f"{seconds} sec") - - -def _credential_values(value: object, *, key: str = "") -> set[str]: - """Extract only secret-shaped values; never include metadata in failures.""" - - found: set[str] = set() - if isinstance(value, dict): - for nested_key, nested_value in value.items(): - found.update(_credential_values(nested_value, key=str(nested_key))) - elif isinstance(value, list): - for nested_value in value: - found.update(_credential_values(nested_value, key=key)) - elif isinstance(value, str) and len(value) >= 8: - normalized_key = key.casefold().replace("_", "").replace("-", "") - if "token" in normalized_key or normalized_key in { - "apikey", - "secret", - "password", - }: - found.add(value) - return found - - -def _load_grok_oauth_secrets() -> set[str]: - from mobilerun.config_manager.credential_paths import GROK_OAUTH_CREDENTIAL_PATH - - credential_path = Path(GROK_OAUTH_CREDENTIAL_PATH) - if not credential_path.is_file(): - pytest.fail( - "Grok OAuth credentials are required; run `mobilerun configure grok` " - "before the opted-in E2E suite" - ) - try: - profiles = json.loads(credential_path.read_text(encoding="utf-8")) - except (OSError, json.JSONDecodeError): - pytest.fail("MobileRun's OAuth credential file is missing or malformed") - - grok_profile = profiles.get("grokOauth") if isinstance(profiles, dict) else None - if not isinstance(grok_profile, dict) or not ( - grok_profile.get("accessToken") or grok_profile.get("refreshToken") - ): - pytest.fail( - "the grokOauth slot is missing; run `mobilerun configure grok` " - "before the opted-in E2E suite" - ) - return _credential_values(grok_profile) - - -def _file_fingerprint(path: Path) -> _FileFingerprint: - """Capture only existence, SHA-256, and mtime; never parse credential data.""" - - try: - before = path.stat() - except FileNotFoundError: - return _FileFingerprint(exists=False, sha256=None, mtime_ns=None) - except OSError: - pytest.fail("could not stat the external Grok CLI credential file") - - digest = hashlib.sha256() - try: - with path.open("rb") as credential_file: - for chunk in iter(lambda: credential_file.read(1024 * 1024), b""): - digest.update(chunk) - after = path.stat() - except OSError: - pytest.fail("could not hash the external Grok CLI credential file") - if before.st_mtime_ns != after.st_mtime_ns: - pytest.fail("the external Grok CLI credential changed while being hashed") - return _FileFingerprint( - exists=True, - sha256=digest.hexdigest(), - mtime_ns=after.st_mtime_ns, - ) - - -def _write_e2e_config( - path: Path, serial: str, trajectory_path: Path -) -> None: - import yaml - - from mobilerun.config_manager import MobileConfig - from mobilerun.config_manager.migrations import CURRENT_VERSION - - config = MobileConfig() - config.agent.max_steps = 30 - config.agent.streaming = False - config.agent.app_cards.enabled = False - config.device.serial = serial - config.device.platform = "android" - config.device.use_tcp = True - config.device.portal_mode = "required" - config.device.auto_setup = False - config.telemetry.enabled = False - config.tracing.enabled = False - config.logging.debug = False - config.logging.rich_text = False - config.logging.save_trajectory = "action" - config.logging.trajectory_gifs = False - config.logging.trajectory_path = str(trajectory_path) - config.mcp.enabled = False - - payload = config.to_dict() - payload["_version"] = CURRENT_VERSION - path.write_text( - yaml.safe_dump(payload, default_flow_style=False, sort_keys=False), - encoding="utf-8", - ) - - -def _portal_version(serial: str) -> str: - output = _adb( - serial, - "shell", - "content", - "query", - "--uri", - "content://com.mobilerun.portal/version", - ) - match = re.search(r"\bresult=(\{.*\})\s*$", output) - if not match: - pytest.fail("Portal content provider did not return version evidence") - try: - payload = json.loads(match.group(1)) - except json.JSONDecodeError: - pytest.fail("Portal content provider returned malformed version evidence") - if payload.get("status") != "success": - pytest.fail("Portal content provider version query was not successful") - version = payload.get("result") or payload.get("data") - if not isinstance(version, str) or not version.strip(): - pytest.fail("Portal content provider returned an empty version") - return version.strip() - - -def _run_public_cli( - repo_root: Path, - serial: str, - *, - args: tuple[str, ...], - secrets: tuple[str, ...], - output_path: Path, - timeout: float, -) -> str: - environment = os.environ.copy() - environment.update( - { - "ANDROID_SERIAL": serial, - "MOBILERUN_TELEMETRY_ENABLED": "false", - "DROIDRUN_TELEMETRY_ENABLED": "false", - } - ) - try: - completed = subprocess.run( - [sys.executable, "-m", "mobilerun", *args], - cwd=repo_root, - env=environment, - check=False, - capture_output=True, - text=True, - timeout=timeout, - ) - except subprocess.TimeoutExpired: - pytest.fail(f"`mobilerun {args[0]}` exceeded its E2E timeout") - output = _ANSI_ESCAPE.sub("", completed.stdout + "\n" + completed.stderr) - _assert_credentials_redacted(output, secrets) - output_path.write_text(output, encoding="utf-8") - assert completed.returncode == 0, _sanitized_tail(output, secrets) - return output - - -def _setup_portal( - repo_root: Path, - serial: str, - *, - secrets: tuple[str, ...], - artifact_dir: Path, -) -> None: - last_output = "" - for attempt in range(1, 4): - output_name = ( - "mobilerun-setup.txt" - if attempt == 1 - else f"mobilerun-setup-retry-{attempt - 1}.txt" - ) - last_output = _run_public_cli( - repo_root, - serial, - args=("setup", "--device", serial), - secrets=secrets, - output_path=artifact_dir / output_name, - timeout=5 * 60, - ) - assert "setup complete!" in last_output.casefold(), ( - "`mobilerun setup` did not report successful Portal setup: " - + _sanitized_tail(last_output, secrets) - ) - assert "setup failed" not in last_output.casefold(), ( - "`mobilerun setup` reported failure: " - + _sanitized_tail(last_output, secrets) - ) - if ( - "did not become responsive" not in last_output.casefold() - and _portal_accessibility_enabled(serial) - ): - return - if attempt < 3: - time.sleep(5.0) - pytest.fail( - "Portal did not become responsive with accessibility enabled after " - "three public `mobilerun setup` attempts: " - + _sanitized_tail(last_output, secrets) - ) - - -def _ping_portal( - repo_root: Path, - serial: str, - *, - use_tcp: bool, - secrets: tuple[str, ...], - artifact_dir: Path, -) -> None: - mode_flag = "--tcp" if use_tcp else "--no-tcp" - mode_name = "tcp" if use_tcp else "content" - output = _run_public_cli( - repo_root, - serial, - args=("ping", "--device", serial, mode_flag, "--no-debug"), - secrets=secrets, - output_path=artifact_dir / f"mobilerun-ping-{mode_name}.txt", - timeout=60, - ) - mode = "TCP" if use_tcp else "content-provider" - assert "portal is installed and accessible" in output.casefold(), ( - f"{mode} Portal ping did not report success: " - + _sanitized_tail(output, secrets) - ) - - -def _doctor_portal( - repo_root: Path, - serial: str, - *, - secrets: tuple[str, ...], - artifact_dir: Path, -) -> None: - output = _run_public_cli( - repo_root, - serial, - args=("doctor", "--device", serial, "--no-debug"), - secrets=secrets, - output_path=artifact_dir / "mobilerun-doctor.txt", - timeout=5 * 60, - ) - assert "mobilerun doctor" in output.casefold(), ( - "`mobilerun doctor` did not start correctly" - ) - assert re.search(r"(?mi)^\s*Portal Version\s{2,}.*$", output), ( - "`mobilerun doctor` did not report its Portal Version check" - ) - required_checks = ( - "Device", - "Portal", - "Accessibility", - "Content Provider", - "State (content)", - "Screenshot (content)", - "TCP Mode", - "State (tcp)", - "Screenshot (tcp)", - ) - # Rich wraps long doctor rows according to the terminal width (for - # example, State (content) can continue on the next line). Bound each - # section by the following row so a checkmark from a later row cannot make - # an earlier failed row pass accidentally. - row_boundaries = (*required_checks, "Keyboard") - for index, check_name in enumerate(required_checks): - next_name = row_boundaries[index + 1] - row = re.search( - rf"(?ms)^\s*{re.escape(check_name)}\s{{2,}}.*?" - rf"(?=^\s*{re.escape(next_name)}\s{{2,}})", - output, - ) - assert row and "โœ“" in row.group(0), ( - f"`mobilerun doctor` did not pass its {check_name} check" - ) - assert not re.search(r"(?mi)^\s*\d+ issue\(s\):", output), ( - "`mobilerun doctor` reported one or more failing checks: " - + _sanitized_tail(output, secrets) - ) - - -@pytest.fixture(scope="session") -def live_context() -> _LiveContext: - serial = os.environ.get("ANDROID_SERIAL") - if serial != _EXPECTED_SERIAL: - pytest.fail( - "the opted-in Grok E2E suite requires " - f"ANDROID_SERIAL={_EXPECTED_SERIAL} exactly" - ) - if shutil.which("adb") is None: - pytest.fail("adb is required for the opted-in Android E2E suite") - - api_key = os.environ.get("XAI_API_KEY", "") - if not api_key: - pytest.fail("XAI_API_KEY is required for the opted-in Grok API E2E cases") - configured_secrets = _load_grok_oauth_secrets() - configured_secrets.add(api_key) - secrets = tuple(configured_secrets) - - repo_root = Path(__file__).resolve().parents[2] - configured_artifact_root = os.environ.get("MOBILERUN_GROK_E2E_ARTIFACTS") - if configured_artifact_root: - artifact_base = Path(configured_artifact_root).expanduser().resolve() - artifact_base.mkdir(parents=True, exist_ok=True) - artifact_root = artifact_base / ( - f"run-{time.strftime('%Y%m%d-%H%M%S')}-{os.getpid()}-{time.time_ns()}" - ) - artifact_root.mkdir() - else: - artifact_root = Path( - tempfile.mkdtemp(prefix="mobilerun-grok-android16-") - ).resolve() - - external_grok_auth = Path.home() / ".grok" / "auth.json" - external_grok_before = _file_fingerprint(external_grok_auth) - emulator = _start_session_emulator(artifact_root) - context = _LiveContext( - serial=serial, - repo_root=repo_root, - artifact_root=artifact_root, - configured_secrets=secrets, - emulator=emulator, - ) - try: - yield context - finally: - _stop_owned_emulator(emulator) - external_grok_after = _file_fingerprint(external_grok_auth) - hash_unchanged = external_grok_before.sha256 == external_grok_after.sha256 - mtime_unchanged = ( - external_grok_before.mtime_ns == external_grok_after.mtime_ns - ) - existence_unchanged = ( - external_grok_before.exists == external_grok_after.exists - ) - (artifact_root / "external-grok-cli-auth-invariant.json").write_text( - json.dumps( - { - "path": "~/.grok/auth.json", - "existence_unchanged": existence_unchanged, - "sha256_unchanged": hash_unchanged, - "mtime_unchanged": mtime_unchanged, - "contents_parsed": False, - }, - indent=2, - ), - encoding="utf-8", - ) - if not (existence_unchanged and hash_unchanged and mtime_unchanged): - pytest.fail( - "MobileRun's live Grok suite modified the separate Grok CLI " - "credential at ~/.grok/auth.json" - ) - - -_ANSI_ESCAPE = re.compile(r"\x1b\[[0-?]*[ -/]*[@-~]") -_TOKEN_ASSIGNMENT = re.compile( - r"(?i)(?:access[_-]?token|refresh[_-]?token|authorization)\s*[:=]\s*" - r"[\"']?(?:bearer\s+)?[A-Za-z0-9._~+/-]{12,}" -) - - -def _assert_credentials_redacted(output: str, secrets: tuple[str, ...]) -> None: - for secret in secrets: - assert secret not in output, "Mobilerun output exposed a configured secret" - assert not _TOKEN_ASSIGNMENT.search(output), ( - "Mobilerun output exposed a token-shaped credential assignment" - ) - - -def _assert_artifacts_redacted(root: Path, secrets: tuple[str, ...]) -> None: - """Scan every artifact as bytes and every UTF-8 artifact as text.""" - - for artifact in sorted(path for path in root.rglob("*") if path.is_file()): - data = artifact.read_bytes() - relative_path = artifact.relative_to(root) - for secret in secrets: - assert secret.encode("utf-8") not in data, ( - f"artifact {relative_path} exposed a configured secret" - ) - try: - text = data.decode("utf-8") - except UnicodeDecodeError: - continue - assert not _TOKEN_ASSIGNMENT.search(text), ( - f"text artifact {relative_path} exposed a token-shaped credential" - ) - - -def _sanitized_tail(output: str, secrets: tuple[str, ...]) -> str: - sanitized = output - for secret in secrets: - sanitized = sanitized.replace(secret, "[REDACTED]") - sanitized = _TOKEN_ASSIGNMENT.sub("credential=[REDACTED]", sanitized) - return sanitized[-4000:] - - -def _jsonable(value: Any) -> Any: - if hasattr(value, "model_dump"): - return value.model_dump(mode="json") - if isinstance(value, dict): - return {str(key): _jsonable(item) for key, item in value.items()} - if isinstance(value, (list, tuple)): - return [_jsonable(item) for item in value] - if isinstance(value, (str, int, float, bool)) or value is None: - return value - return str(value) - - -def _connect_core_device(context: _LiveContext) -> Any: - if Mobilerun is None: - pytest.fail( - "mobilerun_core.Mobilerun is unavailable in the live E2E environment" - ) - try: - device = Mobilerun().connect( - context.serial, - backend="local-android-adb", - portal_mode="required", - ) - except Exception as exc: - pytest.fail(f"public mobilerun-core device connection failed: {exc}") - assert device.capabilities.get("platform") == "android" - return device - - -def _capture_core_evidence( - context: _LiveContext, - artifact_dir: Path, - *, - label: str, - ensure_home: bool, -) -> _CoreEvidence: - """Use only the public mobile-harness surface for control and evidence.""" - - device = _connect_core_device(context) - if ensure_home: - # Snapshot restores may preserve an asleep display. Use only the - # public mobile-harness key/swipe surface for the ordinary wake, - # unlock, and Home sequence before observing the framebuffer. - device.key("wakeup") - time.sleep(0.5) - width, height = device.screen_size() - device.swipe(width // 2, height * 4 // 5, width // 2, height // 5, ms=350) - time.sleep(0.5) - device.key("home") - device.wait_for_idle(timeout=5.0) - current_app_id = device.current_app_id() or "" - ui = device.ui() - screenshot = base64.b64decode(device.screenshot(hide_overlay=True)) - ui_path = artifact_dir / f"{label}-core-ui.json" - screenshot_path = artifact_dir / f"{label}-core.png" - ui_path.write_text( - json.dumps(_jsonable(ui), ensure_ascii=False, indent=2), - encoding="utf-8", - ) - screenshot_path.write_bytes(screenshot) - assert screenshot.startswith(b"\x89PNG\r\n\x1a\n"), ( - "mobilerun-core evidence screenshot was not a PNG" - ) - from PIL import Image - - with Image.open(io.BytesIO(screenshot)) as image: - if image.convert("RGB").getbbox() is None: - raise _AllBlackCoreScreenshot( - "public mobilerun-core returned an all-black screenshot after " - "wake, unlock, and Home" - ) - return _CoreEvidence( - current_app_id=current_app_id, - ui_path=ui_path, - screenshot_path=screenshot_path, - ) - - -def _portal_accessibility_enabled(serial: str) -> bool: - enabled_services = _adb( - serial, - "shell", - "settings", - "get", - "secure", - "enabled_accessibility_services", - ) - return "com.mobilerun.portal" in enabled_services - - -def _setup_and_check_device( - context: _LiveContext, - artifact_dir: Path, -) -> _DevicePreflight: - """Run public setup/ping/doctor and independent readiness assertions.""" - - _setup_portal( - context.repo_root, - context.serial, - secrets=context.configured_secrets, - artifact_dir=artifact_dir, - ) - - android_version = _adb( - context.serial, "shell", "getprop", "ro.build.version.release" - ) - assert android_version == "16", ( - f"{context.serial} must run Android 16 during scenario preflight " - f"(reported {android_version!r})" - ) - sdk_level = _adb( - context.serial, "shell", "getprop", "ro.build.version.sdk" - ) - assert sdk_level == "36", ( - f"{context.serial} must use Android SDK 36 during scenario preflight " - f"(reported {sdk_level!r})" - ) - - portal_package = _adb( - context.serial, "shell", "pm", "path", "com.mobilerun.portal" - ) - assert portal_package.startswith("package:"), ( - "`mobilerun setup` did not install Mobilerun Portal" - ) - portal_version_after_setup = _portal_version(context.serial) - assert re.fullmatch( - r"v?\d+(?:\.\d+){1,3}(?:[-+][A-Za-z0-9.-]+)?", - portal_version_after_setup, - ), "Portal content provider returned an invalid version string" - assert _portal_accessibility_enabled(context.serial), ( - "`mobilerun setup` did not enable Portal accessibility" - ) - - _ping_portal( - context.repo_root, - context.serial, - use_tcp=False, - secrets=context.configured_secrets, - artifact_dir=artifact_dir, - ) - _ping_portal( - context.repo_root, - context.serial, - use_tcp=True, - secrets=context.configured_secrets, - artifact_dir=artifact_dir, - ) - _doctor_portal( - context.repo_root, - context.serial, - secrets=context.configured_secrets, - artifact_dir=artifact_dir, - ) - - portal_version_after_doctor = _portal_version(context.serial) - assert re.fullmatch( - r"v?\d+(?:\.\d+){1,3}(?:[-+][A-Za-z0-9.-]+)?", - portal_version_after_doctor, - ), "Portal content provider returned an invalid post-doctor version string" - assert _portal_accessibility_enabled(context.serial), ( - "Portal accessibility was not enabled after `mobilerun doctor`" - ) - return _DevicePreflight( - android_version=android_version, - sdk_level=sdk_level, - portal_version_after_setup=portal_version_after_setup, - portal_version_after_doctor=portal_version_after_doctor, - ) - - -def _prepare_scenario_device( - context: _LiveContext, - scenario: _Scenario, - artifact_dir: Path, -) -> _ScenarioDeviceState: - """Restore, boot, set up, and fully validate the device for one scenario.""" - - initial_reset_mode = _restore_clean_snapshot( - context.emulator, context.serial, artifact_dir - ) - final_reset_mode = initial_reset_mode - graphics_fallback_reason: str | None = None - preflight_artifact_dir = artifact_dir - preflight = _setup_and_check_device(context, preflight_artifact_dir) - try: - pre_evidence = _capture_core_evidence( - context, - preflight_artifact_dir, - label="pre-task", - ensure_home=True, - ) - except _AllBlackCoreScreenshot as exc: - if initial_reset_mode != "snapshot": - raise - graphics_fallback_reason = str(exc) - fallback_dir = artifact_dir / "snapshot-graphics-cold-boot" - fallback_dir.mkdir() - (artifact_dir / "snapshot-graphics-fallback.json").write_text( - json.dumps( - { - "snapshot_attempted": True, - "snapshot_restored": True, - "failure": "all_black_public_core_framebuffer", - "reason": graphics_fallback_reason, - "fallback": "cold-boot", - "wipe_data": False, - "rerun_after_fallback": [ - "mobilerun setup", - "mobilerun ping --no-tcp", - "mobilerun ping --tcp", - "mobilerun doctor", - "public mobilerun-core evidence", - ], - }, - indent=2, - ), - encoding="utf-8", - ) - try: - _cold_boot_owned_emulator(context.emulator, fallback_dir) - except RuntimeError as cold_error: - pytest.fail( - "snapshot framebuffer was all-black and the no-wipe-data " - f"cold-boot fallback failed: {cold_error}" - ) - preflight_artifact_dir = fallback_dir - preflight = _setup_and_check_device(context, preflight_artifact_dir) - try: - pre_evidence = _capture_core_evidence( - context, - preflight_artifact_dir, - label="pre-task", - ensure_home=True, - ) - except _AllBlackCoreScreenshot: - pytest.fail( - "public mobilerun-core framebuffer remained all-black after " - "the no-wipe-data cold-boot fallback" - ) - final_reset_mode = "snapshot_graphics_cold_boot_fallback" - - home_package = pre_evidence.current_app_id - assert home_package, "mobilerun-core could not identify the Home package" - timeout_labels = _timeout_labels( - _adb( - context.serial, - "shell", - "settings", - "get", - "system", - "screen_off_timeout", - ) - ) - (artifact_dir / "preflight.json").write_text( - json.dumps( - { - "scenario_id": scenario.scenario_id, - "serial": context.serial, - "avd": _AVD_NAME, - "emulator_port": _EMULATOR_PORT, - "snapshot": _SNAPSHOT_NAME, - "snapshot_attempted": True, - "snapshot_restored": initial_reset_mode == "snapshot", - "initial_reset_mode": initial_reset_mode, - "reset_mode": final_reset_mode, - "snapshot_graphics_fallback_reason": graphics_fallback_reason, - "cold_boot_without_wipe_data": final_reset_mode - in { - "cold_boot_fallback", - "snapshot_graphics_cold_boot_fallback", - }, - "boot_completed": True, - "mobilerun_setup": "passed", - "android_version": preflight.android_version, - "sdk_level": preflight.sdk_level, - "portal_version_after_setup": preflight.portal_version_after_setup, - "portal_version_after_doctor": preflight.portal_version_after_doctor, - "portal_accessibility": "enabled", - "content_provider_ping": "passed", - "tcp_ping": "passed", - "mobilerun_doctor": "passed", - "portal_mode": "required", - "auto_setup_during_task": False, - }, - indent=2, - ), - encoding="utf-8", - ) - _assert_artifacts_redacted(artifact_dir, context.configured_secrets) - return _ScenarioDeviceState( - home_package=home_package, - android_version=preflight.android_version, - timeout_labels=timeout_labels, - pre_evidence=pre_evidence, - ) - - -def _load_and_assert_trajectory(trajectory_root: Path) -> tuple[Path, tuple[dict[str, object], ...]]: - trajectory_files = sorted(trajectory_root.glob("*/trajectory.json")) - assert len(trajectory_files) == 1, ( - "each Grok scenario must produce exactly one trajectory.json artifact" - ) - trajectory_file = trajectory_files[0] - try: - raw_events = json.loads(trajectory_file.read_text(encoding="utf-8")) - except json.JSONDecodeError: - pytest.fail("trajectory.json was not valid JSON") - assert isinstance(raw_events, list) and raw_events, ( - "trajectory.json must contain recorded MobileRun events" - ) - events = tuple(event for event in raw_events if isinstance(event, dict)) - assert len(events) == len(raw_events), "trajectory events must be JSON objects" - - response_events = [ - event - for event in events - if str(event.get("type", "")).endswith("ResponseEvent") - ] - assert any( - isinstance(event.get("usage"), dict) - and int(event["usage"].get("requests", 0) or 0) > 0 - and int(event["usage"].get("total_tokens", 0) or 0) > 0 - for event in response_events - ), "trajectory must include positive request and token usage on a ResponseEvent" - - assert any( - event.get("type") == "ToolExecutionEvent" and event.get("success") is True - for event in events - ), "trajectory must include a successful ToolExecutionEvent" - direct_completion = any( - event.get("type") == "ToolExecutionEvent" - and event.get("tool_name") == "complete" - and event.get("success") is True - and isinstance(event.get("tool_args"), dict) - and event["tool_args"].get("success") is True - for event in events - ) - reasoning_completion = any( - event.get("type") == "ManagerResponseEvent" - and isinstance(event.get("response"), str) - and re.search( - r"]*\bsuccess=[\"']true[\"']", - event["response"], - flags=re.IGNORECASE, - ) - for event in events - ) - assert direct_completion or reasoning_completion, ( - "trajectory must include a successful MobileRun completion result" - ) - - trajectory_dir = trajectory_file.parent - ui_state_files = sorted((trajectory_dir / "ui_states").glob("*.json")) - screenshot_files = sorted((trajectory_dir / "screenshots").glob("*.png")) - assert ui_state_files, "trajectory must include recorded UI-state artifacts" - assert screenshot_files, "trajectory must include screenshot artifacts" - for ui_state_file in ui_state_files: - try: - ui_state = json.loads(ui_state_file.read_text(encoding="utf-8")) - except json.JSONDecodeError: - pytest.fail("a recorded UI-state artifact was not valid JSON") - assert isinstance(ui_state, list), "recorded UI state must be a JSON list" - for screenshot_file in screenshot_files: - screenshot = screenshot_file.read_bytes() - assert screenshot.startswith(b"\x89PNG\r\n\x1a\n"), ( - "recorded screenshot must be a non-empty PNG artifact" - ) - return trajectory_dir, events - - -def _recorded_screen_timeout_label(trajectory_dir: Path) -> str: - """Read the visible Settings summary from independently recorded Portal UI.""" - - for ui_state_file in sorted((trajectory_dir / "ui_states").glob("*.json")): - raw_nodes = json.loads(ui_state_file.read_text(encoding="utf-8")) - if not isinstance(raw_nodes, list): - continue - for index, node in enumerate(raw_nodes): - if not isinstance(node, dict) or str(node.get("text", "")).casefold() != "screen timeout": - continue - for summary in raw_nodes[index + 1 : index + 6]: - if not isinstance(summary, dict): - continue - resource_id = str(summary.get("resourceId", "")) - label = str(summary.get("text", "")).strip() - if resource_id.endswith("id/summary") and label: - return label - pytest.fail("recorded Portal UI did not contain the visible Screen timeout summary") - - -def _run_scenario( - context: _LiveContext, scenario: _Scenario -) -> _ScenarioResult: - artifact_dir = context.artifact_root / scenario.scenario_id - artifact_dir.mkdir(parents=False, exist_ok=False) - trajectory_root = artifact_dir / "trajectories" - config_path = artifact_dir / "config.yaml" - (artifact_dir / "scenario.json").write_text( - json.dumps( - { - "scenario_id": scenario.scenario_id, - "snapshot": _SNAPSHOT_NAME, - "provider": scenario.provider, - "auth_mode": scenario.auth_mode, - "model": "grok-4.5", - "mode_flags": list(scenario.mode_flags), - "portal_mode": "required", - "auto_setup": False, - "tcp": True, - "tracing": False, - "telemetry": False, - "trajectory": "action", - }, - indent=2, - ), - encoding="utf-8", - ) - device_state = _prepare_scenario_device(context, scenario, artifact_dir) - _write_e2e_config(config_path, context.serial, trajectory_root) - command = [ - sys.executable, - "-m", - "mobilerun", - "run", - scenario.task, - "--config", - str(config_path), - "--device", - context.serial, - "--tcp", - "--provider", - scenario.provider, - "--model", - "grok-4.5", - "--steps", - "30", - "--no-stream", - "--no-tracing", - "--no-debug", - "--save-trajectory", - "action", - *scenario.mode_flags, - ] - environment = os.environ.copy() - environment.update( - { - "ANDROID_SERIAL": context.serial, - "MOBILERUN_CONFIG": str(config_path), - "MOBILERUN_TELEMETRY_ENABLED": "false", - "DROIDRUN_TELEMETRY_ENABLED": "false", - } - ) - # The natural-language agent CLI is the Grok integration under test. All - # harness-owned device actions and evidence use mobilerun_core.Mobilerun. - completed = subprocess.run( - command, - cwd=context.repo_root, - env=environment, - check=False, - capture_output=True, - text=True, - timeout=20 * 60, - ) - output = _ANSI_ESCAPE.sub("", completed.stdout + "\n" + completed.stderr) - # OAuth may rotate tokens during inference. Scan for both the credentials - # that existed before the suite and the newly persisted credential values. - runtime_secrets = set(context.configured_secrets) - runtime_secrets.update(_load_grok_oauth_secrets()) - runtime_secrets.add(os.environ["XAI_API_KEY"]) - secrets = tuple(runtime_secrets) - _assert_credentials_redacted(output, secrets) - (artifact_dir / "mobilerun-run.txt").write_text(output, encoding="utf-8") - post_evidence = _capture_core_evidence( - context, - artifact_dir, - label="post-task", - ensure_home=False, - ) - _assert_artifacts_redacted(artifact_dir, secrets) - assert completed.returncode == 0, _sanitized_tail( - output, secrets - ) - trajectory_dir, events = _load_and_assert_trajectory(trajectory_root) - return _ScenarioResult( - output=output, - artifact_dir=artifact_dir, - trajectory_dir=trajectory_dir, - events=events, - device_state=device_state, - post_evidence=post_evidence, - ) - - -def _assert_scenario_state( - context: _LiveContext, - scenario: _Scenario, - result: _ScenarioResult, -) -> None: - normalized_output = " ".join(result.output.casefold().split()) - assert result.post_evidence.current_app_id == result.device_state.home_package, ( - "Grok did not finish the task on the Home screen" - ) - assert _focused_package(context.serial) == result.device_state.home_package, ( - "independent foreground diagnostics did not confirm the Home screen" - ) - if scenario.expected_state == "screen_timeout": - visible_label = _recorded_screen_timeout_label(result.trajectory_dir) - assert " ".join(visible_label.casefold().split()) in normalized_output, ( - "Grok did not report the emulator's visible Screen timeout value" - ) - return - - assert result.device_state.android_version.casefold() in normalized_output, ( - "Grok did not report the emulator's Android version" - ) - assert ( - _adb(context.serial, "shell", "getprop", "ro.build.version.release") - == result.device_state.android_version - ) - - -def test_xai_api_direct_ui_tree_screen_timeout(live_context: _LiveContext) -> None: - result = _run_scenario(live_context, _API_DIRECT) - _assert_scenario_state(live_context, _API_DIRECT, result) - - -def test_xai_api_reasoning_vision_only_android_version( - live_context: _LiveContext, -) -> None: - result = _run_scenario(live_context, _API_REASONING_VISION) - _assert_scenario_state(live_context, _API_REASONING_VISION, result) - - -def test_grok_oauth_direct_ui_tree_screen_timeout( - live_context: _LiveContext, -) -> None: - result = _run_scenario(live_context, _OAUTH_DIRECT) - _assert_scenario_state(live_context, _OAUTH_DIRECT, result) - - -def test_grok_oauth_reasoning_vision_only_android_version( - live_context: _LiveContext, -) -> None: - result = _run_scenario(live_context, _OAUTH_REASONING_VISION) - _assert_scenario_state(live_context, _OAUTH_REASONING_VISION, result) diff --git a/tests/test_grok_android_e2e_helpers.py b/tests/test_grok_android_e2e_helpers.py deleted file mode 100644 index d6926c2a..00000000 --- a/tests/test_grok_android_e2e_helpers.py +++ /dev/null @@ -1,174 +0,0 @@ -from __future__ import annotations - -import base64 -import io -import json -from pathlib import Path -from typing import Any - -import pytest -from PIL import Image - -from tests.e2e import test_grok_android16 as android_e2e - - -class _BlackScreenshotDevice: - capabilities = {"platform": "android"} - - def __init__(self, screenshot: bytes) -> None: - self._screenshot = screenshot - self.actions: list[tuple[Any, ...]] = [] - - def key(self, name: str) -> None: - self.actions.append(("key", name)) - - def screen_size(self) -> tuple[int, int]: - return 1080, 2400 - - def swipe(self, *args: Any, **kwargs: Any) -> None: - self.actions.append(("swipe", *args, kwargs)) - - def wait_for_idle(self, *, timeout: float) -> bool: - self.actions.append(("wait_for_idle", timeout)) - return True - - def current_app_id(self) -> str: - return "com.android.launcher3" - - def ui(self) -> dict[str, object]: - return {"phone_state": {"package_name": "com.android.launcher3"}} - - def screenshot(self, *, hide_overlay: bool) -> str: - assert hide_overlay is True - return base64.b64encode(self._screenshot).decode("ascii") - - -def _black_png() -> bytes: - output = io.BytesIO() - Image.new("RGB", (4, 4), color="black").save(output, format="PNG") - return output.getvalue() - - -def _context(tmp_path: Path) -> android_e2e._LiveContext: - emulator = android_e2e._OwnedEmulator( - binary=tmp_path / "emulator", - artifact_root=tmp_path, - ) - return android_e2e._LiveContext( - serial="emulator-5558", - repo_root=tmp_path, - artifact_root=tmp_path, - configured_secrets=(), - emulator=emulator, - ) - - -def test_public_core_black_png_raises_graphics_readiness_error( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - device = _BlackScreenshotDevice(_black_png()) - monkeypatch.setattr(android_e2e, "_connect_core_device", lambda context: device) - monkeypatch.setattr(android_e2e.time, "sleep", lambda seconds: None) - - with pytest.raises( - android_e2e._AllBlackCoreScreenshot, - match="all-black screenshot", - ): - android_e2e._capture_core_evidence( - _context(tmp_path), - tmp_path, - label="pre-task", - ensure_home=True, - ) - - assert (tmp_path / "pre-task-core.png").read_bytes() == device._screenshot - assert device.actions[0] == ("key", "wakeup") - assert ("key", "home") in device.actions - - -def test_snapshot_black_framebuffer_uses_no_wipe_cold_boot_and_rechecks( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - artifact_dir = tmp_path / "scenario" - artifact_dir.mkdir() - context = _context(tmp_path) - setup_dirs: list[Path] = [] - capture_dirs: list[Path] = [] - cold_boot_dirs: list[Path] = [] - preflight = android_e2e._DevicePreflight( - android_version="16", - sdk_level="36", - portal_version_after_setup="0.7.1", - portal_version_after_doctor="0.7.1", - ) - recovered = android_e2e._CoreEvidence( - current_app_id="com.android.launcher3", - ui_path=artifact_dir / "recovered-ui.json", - screenshot_path=artifact_dir / "recovered.png", - ) - - monkeypatch.setattr( - android_e2e, - "_restore_clean_snapshot", - lambda emulator, serial, artifacts: "snapshot", - ) - - def setup_and_check( - unused_context: android_e2e._LiveContext, - artifacts: Path, - ) -> android_e2e._DevicePreflight: - setup_dirs.append(artifacts) - return preflight - - monkeypatch.setattr(android_e2e, "_setup_and_check_device", setup_and_check) - - def capture( - unused_context: android_e2e._LiveContext, - artifacts: Path, - *, - label: str, - ensure_home: bool, - ) -> android_e2e._CoreEvidence: - assert label == "pre-task" - assert ensure_home is True - capture_dirs.append(artifacts) - if len(capture_dirs) == 1: - raise android_e2e._AllBlackCoreScreenshot("all-black public frame") - return recovered - - monkeypatch.setattr(android_e2e, "_capture_core_evidence", capture) - monkeypatch.setattr( - android_e2e, - "_cold_boot_owned_emulator", - lambda emulator, artifacts: cold_boot_dirs.append(artifacts), - ) - monkeypatch.setattr(android_e2e, "_adb", lambda serial, *args: "30000") - monkeypatch.setattr( - android_e2e, - "_timeout_labels", - lambda timeout_ms: ("30 seconds",), - ) - - result = android_e2e._prepare_scenario_device( - context, - android_e2e._API_DIRECT, - artifact_dir, - ) - - fallback_dir = artifact_dir / "snapshot-graphics-cold-boot" - assert setup_dirs == [artifact_dir, fallback_dir] - assert capture_dirs == [artifact_dir, fallback_dir] - assert cold_boot_dirs == [fallback_dir] - assert result.pre_evidence is recovered - fallback = json.loads( - (artifact_dir / "snapshot-graphics-fallback.json").read_text() - ) - assert fallback["failure"] == "all_black_public_core_framebuffer" - assert fallback["wipe_data"] is False - summary = json.loads((artifact_dir / "preflight.json").read_text()) - assert summary["snapshot_attempted"] is True - assert summary["snapshot_restored"] is True - assert summary["reset_mode"] == "snapshot_graphics_cold_boot_fallback" - assert summary["cold_boot_without_wipe_data"] is True diff --git a/tests/test_grok_api.py b/tests/test_grok_api.py index 6b009513..01dd95f5 100644 --- a/tests/test_grok_api.py +++ b/tests/test_grok_api.py @@ -79,12 +79,12 @@ def _xai_usage() -> ResponseUsage: def test_grok_api_key_variant_is_first_class_xai_responses_provider() -> None: - variant = resolve_provider_variant("grok", "api_key") + 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("grok", "api_key") == ("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" @@ -92,10 +92,10 @@ def test_grok_api_key_variant_is_first_class_xai_responses_provider() -> None: def test_grok_oauth_variant_shares_the_canonical_model_catalog() -> None: - variant = resolve_provider_variant("grok", "oauth") + variant = resolve_provider_variant("xai", "oauth") - assert variant.id == "grok_oauth" - assert variant.runtime_provider_name == "grok_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 @@ -109,21 +109,27 @@ def test_grok_oauth_variant_shares_the_canonical_model_catalog() -> None: def test_grok_model_aliases_normalize_to_canonical_id( auth_mode: str, model_alias: str ) -> None: - assert normalize_model_id_for_variant("grok", auth_mode, model_alias) == "grok-4.5" + assert normalize_model_id_for_variant("xai", auth_mode, model_alias) == "grok-4.5" -@pytest.mark.parametrize("alias", ("grok", "xai", "x.ai", "XAI")) +@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("grok", "api_key") + variant = resolve_provider_variant("xai", "api_key") profile = create_profile_for_variant( variant, SetupSelection( - family_id="grok", + family_id="xai", variant_id="XAI", auth_mode="api_key", model="grok-build-latest", @@ -133,7 +139,7 @@ def test_grok_profile_wires_api_base_context_and_environment_key(monkeypatch) -> ) assert profile.provider == "XAI" - assert profile.provider_family == "grok" + assert profile.provider_family == "xai" assert profile.model == "grok-4.5" assert profile.temperature == 0.4 assert profile.base_url == XAI_API_BASE @@ -153,11 +159,11 @@ def test_grok_profile_resolves_saved_api_key(monkeypatch, tmp_path) -> None: ) monkeypatch.setattr(env_keys, "AUTH_PROFILES_PATH", credential_path) monkeypatch.delenv("XAI_API_KEY", raising=False) - variant = resolve_provider_variant("grok", "api_key") + variant = resolve_provider_variant("xai", "api_key") profile = create_profile_for_variant( variant, SetupSelection( - family_id="grok", + family_id="xai", variant_id="XAI", auth_mode="api_key", model="grok-4.5", @@ -507,12 +513,12 @@ async def parse_async(**kwargs: Any) -> Any: def test_xai_loader_uses_environment_key_for_direct_runtime(monkeypatch) -> None: monkeypatch.setenv("XAI_API_KEY", "xai-runtime-key") - llm = load_llm("grok", model="grok-4.5") + llm = load_llm("xai", model="grok-4.5") assert llm.api_key == "xai-runtime-key" -@pytest.mark.parametrize("alias", ("grok", "xai", "x.ai", "XAI")) +@pytest.mark.parametrize("alias", ("xai", "XAI")) def test_xai_runtime_aliases_default_to_canonical_model( alias: str, monkeypatch ) -> None: @@ -523,6 +529,21 @@ def test_xai_runtime_aliases_default_to_canonical_model( 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") diff --git a/tests/test_grok_cli.py b/tests/test_grok_cli.py index aed93df8..8ce8ab76 100644 --- a/tests/test_grok_cli.py +++ b/tests/test_grok_cli.py @@ -1,8 +1,10 @@ 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 @@ -10,7 +12,7 @@ from mobilerun.config_manager import MobileConfig -def test_grok_oauth_credentials_are_detected_by_nested_slot(tmp_path) -> None: +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( @@ -23,7 +25,7 @@ def test_grok_oauth_credentials_are_detected_by_nested_slot(tmp_path) -> None: ) assert configure_wizard._oauth_credentials_present( - str(credential_path), "grok_oauth" + str(credential_path), "xai_oauth" ) credential_path.write_text( @@ -31,11 +33,11 @@ def test_grok_oauth_credentials_are_detected_by_nested_slot(tmp_path) -> None: encoding="utf-8", ) assert not configure_wizard._oauth_credentials_present( - str(credential_path), "grok_oauth" + str(credential_path), "xai_oauth" ) -def test_wizard_prepares_grok_oauth_with_selected_model(tmp_path) -> None: +def test_wizard_prepares_xai_oauth_with_selected_model(tmp_path) -> None: calls: list[dict] = [] callbacks = ConfigureWizardCallbacks( run_openai_oauth_login=lambda **kwargs: None, @@ -46,7 +48,7 @@ def test_wizard_prepares_grok_oauth_with_selected_model(tmp_path) -> None: configure_wizard._prepare_variant_auth( callbacks=callbacks, - variant=SimpleNamespace(id="grok_oauth"), + variant=SimpleNamespace(id="xai_oauth"), credential_path=str(tmp_path / "auth-profiles.json"), selected_model="grok-4.5", ) @@ -59,7 +61,7 @@ def test_wizard_prepares_grok_oauth_with_selected_model(tmp_path) -> None: ] -def test_configure_grok_command_forwards_device_code_options( +def test_configure_xai_command_forwards_device_code_options( monkeypatch, tmp_path ) -> None: calls: list[dict] = [] @@ -74,7 +76,7 @@ def test_configure_grok_command_forwards_device_code_options( cli_main.cli, [ "configure", - "grok", + "xai", "--credential-path", str(credential_path), "--model", @@ -98,24 +100,33 @@ def test_configure_grok_command_forwards_device_code_options( ] -def test_configure_help_advertises_grok_provider_and_login_command() -> None: +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"]) - grok_help = runner.invoke(cli_main.cli, ["configure", "grok", "--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 "grok" in configure_help.output.lower() - assert grok_help.exit_code == 0 - assert "--device-code" in grok_help.output - assert "native xAI OAuth" in grok_help.output + 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", "grok_oauth")), + (("api_key", "XAI"), ("oauth", "xai_oauth")), ) -def test_exact_grok_configure_forms_keep_provider_and_auth_fixed( +def test_exact_xai_configure_forms_keep_provider_and_auth_fixed( monkeypatch, auth_mode: str, expected_provider: str ) -> None: config = MobileConfig() @@ -154,7 +165,7 @@ def choose_model(models, *, default_model, allow_back=True): # type: ignore[no- configure_wizard, "select_prompt", lambda *args, **kwargs: pytest.fail( - "fixed Grok configure flow unexpectedly reopened the top-level menu" + "fixed XAI configure flow unexpectedly reopened the top-level menu" ), ) monkeypatch.setattr( @@ -165,7 +176,7 @@ def choose_model(models, *, default_model, allow_back=True): # type: ignore[no- result = CliRunner().invoke( cli_main.cli, - ["configure", "--provider", "grok", "--auth-mode", auth_mode], + ["configure", "--provider", "XAI", "--auth-mode", auth_mode], ) assert result.exit_code == 0, result.output @@ -175,4 +186,94 @@ def choose_model(models, *, default_model, allow_back=True): # type: ignore[no- assert { (profile.provider, profile.provider_family, profile.auth_mode, profile.model) for profile in config.llm_profiles.values() - } == {(expected_provider, "grok", auth_mode, "grok-4.5")} + } == {(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 index 526394ef..2df1f201 100644 --- a/tests/test_grok_oauth.py +++ b/tests/test_grok_oauth.py @@ -43,7 +43,7 @@ GrokOAuthSessionManager, _parse_callback_query, ) -from mobilerun.config_manager import env_keys +from mobilerun.config_manager import auth_profile_store, env_keys from mobilerun.config_manager.auth_profile_store import ( AuthProfileFormatError, AuthProfileStore, @@ -176,6 +176,16 @@ def test_auth_profile_store_preserves_siblings_and_writes_private_file(tmp_path: 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") @@ -1357,7 +1367,7 @@ def test_oauth_responses_adapter_pins_proxy_and_omits_controls(tmp_path: Path): ): assert key not in model_kwargs - with pytest.raises(ValueError, match="not supported with Grok OAuth"): + with pytest.raises(ValueError, match="not supported with XAI OAuth"): GrokOAuth( model="other-model", oauth_credential_path=str(tmp_path / "other.json"), @@ -1368,7 +1378,7 @@ def test_oauth_responses_adapter_pins_proxy_and_omits_controls(tmp_path: Path): oauth_credential_path=str(tmp_path / "missing.json"), ) assert no_fallback.api_key == "oauth" - with pytest.raises(ValueError, match="No Grok OAuth credentials"): + with pytest.raises(ValueError, match="No XAI OAuth credentials"): no_fallback._oauth_manager.get_valid_credentials() diff --git a/tests/test_llm_picker.py b/tests/test_llm_picker.py index 0cc91415..93ef3c64 100644 --- a/tests/test_llm_picker.py +++ b/tests/test_llm_picker.py @@ -125,7 +125,6 @@ async def parse_async(**kwargs: Any) -> Any: for payload in (sync_payload, async_payload): assert {"temperature", "top_p"}.isdisjoint(payload) assert payload["max_output_tokens"] == 32 - assert payload["tool_choice"] == "none" @pytest.mark.parametrize( diff --git a/tests/test_usage.py b/tests/test_usage.py index df43aaff..c8683b25 100644 --- a/tests/test_usage.py +++ b/tests/test_usage.py @@ -127,7 +127,7 @@ def test_openai_responses_extracts_completed_stream_additional_usage_fallback() }, ) - usage = get_usage_from_response("GrokOAuth", chat_response) + usage = get_usage_from_response("xai_oauth", chat_response) assert usage.request_tokens == 13 assert usage.response_tokens == 6 From a4f28c8633a66697f77a52dd30738191373034d4 Mon Sep 17 00:00:00 2001 From: "rasul.osmanbayli" Date: Thu, 13 Aug 2026 12:54:56 +0400 Subject: [PATCH 3/4] Format XAI provider changes --- mobilerun/agent/utils/llm_picker.py | 12 ++--- .../agent/utils/oauth/anthropic_oauth_llm.py | 23 +++++---- .../oauth/gemini_oauth_code_assist_llm.py | 23 +++++---- mobilerun/agent/utils/oauth/grok_oauth_llm.py | 43 ++++++----------- mobilerun/cli/configure_wizard.py | 4 +- mobilerun/cli/oauth_actions.py | 15 +++--- tests/test_grok_api.py | 4 +- tests/test_grok_oauth.py | 48 ++++++++----------- 8 files changed, 75 insertions(+), 97 deletions(-) diff --git a/mobilerun/agent/utils/llm_picker.py b/mobilerun/agent/utils/llm_picker.py index 94fc4d4c..2aa92fe3 100644 --- a/mobilerun/agent/utils/llm_picker.py +++ b/mobilerun/agent/utils/llm_picker.py @@ -228,9 +228,7 @@ def _get_model_kwargs(self, **kwargs: Any) -> dict[str, Any]: 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 - ) + 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. @@ -259,9 +257,7 @@ def structured_predict( input=message_dicts, text_format=output_cls, store=self.store, - **self._sanitize_structured_call_kwargs( - dict(llm_kwargs or {}) - ), + **self._sanitize_structured_call_kwargs(dict(llm_kwargs or {})), ) if response.output_parsed is not None: return response.output_parsed @@ -294,9 +290,7 @@ async def astructured_predict( input=message_dicts, text_format=output_cls, store=self.store, - **self._sanitize_structured_call_kwargs( - dict(llm_kwargs or {}) - ), + **self._sanitize_structured_call_kwargs(dict(llm_kwargs or {})), ) if response.output_parsed is not None: return response.output_parsed diff --git a/mobilerun/agent/utils/oauth/anthropic_oauth_llm.py b/mobilerun/agent/utils/oauth/anthropic_oauth_llm.py index 0cfca70a..29de8e06 100644 --- a/mobilerun/agent/utils/oauth/anthropic_oauth_llm.py +++ b/mobilerun/agent/utils/oauth/anthropic_oauth_llm.py @@ -267,16 +267,19 @@ def _persist_credentials(self) -> None: if not self.credential_path: return path = Path(self.credential_path).expanduser() - 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(), - }) + 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(), + }, + ) def _token_headers(self) -> Dict[str, str]: headers = { 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 f5efd952..455b0a70 100644 --- a/mobilerun/agent/utils/oauth/gemini_oauth_code_assist_llm.py +++ b/mobilerun/agent/utils/oauth/gemini_oauth_code_assist_llm.py @@ -297,16 +297,19 @@ def _persist_credentials(self) -> None: return path = Path(self.credential_path).expanduser() - 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 - ), - }) + 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 + ), + }, + ) def _metadata_payload(self) -> Dict[str, str]: return { diff --git a/mobilerun/agent/utils/oauth/grok_oauth_llm.py b/mobilerun/agent/utils/oauth/grok_oauth_llm.py index c7cbc2d7..30aa7ccb 100644 --- a/mobilerun/agent/utils/oauth/grok_oauth_llm.py +++ b/mobilerun/agent/utils/oauth/grok_oauth_llm.py @@ -119,9 +119,7 @@ def _parse_callback_query(query: str) -> dict[str, str | None]: 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 + 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, @@ -143,7 +141,9 @@ class GrokOAuthReloginRequired(GrokOAuthError): def _safe_error_code(value: object) -> str | None: - return value if isinstance(value, str) and _OAUTH_ERROR_CODE.fullmatch(value) else None + return ( + value if isinstance(value, str) and _OAUTH_ERROR_CODE.fullmatch(value) else None + ) def _safe_token_error_code(value: object) -> str | None: @@ -229,9 +229,7 @@ def is_valid(self, *, skew_ms: int = DEFAULT_REFRESH_SKEW_SECONDS * 1000) -> boo class GrokOAuthCredentialStore: - def __init__( - self, path: str | Path = DEFAULT_GROK_OAUTH_CREDENTIAL_PATH - ) -> None: + def __init__(self, path: str | Path = DEFAULT_GROK_OAUTH_CREDENTIAL_PATH) -> None: self.path = Path(path).expanduser() self.profile_store = AuthProfileStore(self.path) @@ -251,9 +249,7 @@ 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_client = jwks_client or jwt.PyJWKClient(DEFAULT_GROK_OAUTH_JWKS_URL) def validate(self, token: str, *, nonce: str | None) -> dict[str, Any]: signing_key = self._jwks_client.get_signing_key_from_jwt(token).key @@ -341,9 +337,7 @@ def _credentials_from_token_response( 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 - ): + 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." ) @@ -444,9 +438,7 @@ def exchange_authorization_code( "code_verifier": code_verifier, } ) - credentials = self._credentials_from_token_response( - payload, nonce=nonce - ) + credentials = self._credentials_from_token_response(payload, nonce=nonce) self.set_initial_credentials(credentials) return credentials @@ -535,9 +527,7 @@ def _authorize( 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]: + def sync_auth_flow(self, request: httpx.Request) -> Iterator[httpx.Request]: credentials = self.manager.get_valid_credentials() self._authorize(request, credentials) response = yield request @@ -607,11 +597,7 @@ def __init__( 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 - ), + expires_at_ms=(oauth_expires_at_ms if oauth_access_token else 0), ) ) @@ -887,7 +873,10 @@ def _login_device_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)): + 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))) @@ -920,9 +909,7 @@ def _login_device_code( token_payload = _safe_json_object( token_response, context="xAI device token response" ) - credentials = manager._credentials_from_token_response( - token_payload - ) + credentials = manager._credentials_from_token_response(token_payload) manager.set_initial_credentials(credentials) return credentials try: diff --git a/mobilerun/cli/configure_wizard.py b/mobilerun/cli/configure_wizard.py index 8afe3da2..115cc9b5 100644 --- a/mobilerun/cli/configure_wizard.py +++ b/mobilerun/cli/configure_wizard.py @@ -557,9 +557,7 @@ 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 - ): + 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: diff --git a/mobilerun/cli/oauth_actions.py b/mobilerun/cli/oauth_actions.py index 1ce171c4..2e29ad59 100644 --- a/mobilerun/cli/oauth_actions.py +++ b/mobilerun/cli/oauth_actions.py @@ -124,12 +124,15 @@ def run_anthropic_setup_token_oauth( def save_anthropic_setup_token(credential_path: str, token: str) -> None: - AuthProfileStore(credential_path).update_slot("claudeAiOauth", { - "accessToken": token, - "refreshToken": None, - "expiresAt": None, - "scopes": [], - }) + AuthProfileStore(credential_path).update_slot( + "claudeAiOauth", + { + "accessToken": token, + "refreshToken": None, + "expiresAt": None, + "scopes": [], + }, + ) def run_anthropic_oauth_setup(credential_path: str) -> None: diff --git a/tests/test_grok_api.py b/tests/test_grok_api.py index 01dd95f5..6e946156 100644 --- a/tests/test_grok_api.py +++ b/tests/test_grok_api.py @@ -409,9 +409,7 @@ async def create_async(**kwargs: Any): # type: ignore[no-untyped-def] async def collect_async(): # type: ignore[no-untyped-def] return [ item - async for item in await llm.astream_chat( - messages, **dict(runtime_kwargs) - ) + async for item in await llm.astream_chat(messages, **dict(runtime_kwargs)) ] async_result = asyncio.run(collect_async())[-1] diff --git a/tests/test_grok_oauth.py b/tests/test_grok_oauth.py index 2df1f201..8f9edcdc 100644 --- a/tests/test_grok_oauth.py +++ b/tests/test_grok_oauth.py @@ -130,8 +130,7 @@ def _responses_sse() -> bytes: ) return ( "".join( - f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" - for event in events + f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in events ) + "data: [DONE]\n\n" ).encode() @@ -176,7 +175,9 @@ def test_auth_profile_store_preserves_siblings_and_writes_private_file(tmp_path: assert not list(tmp_path.glob(".auth-profiles.json.*.tmp")) -def test_auth_profile_store_writes_when_fchmod_is_unavailable(monkeypatch, tmp_path: Path): +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) @@ -659,7 +660,9 @@ def handler(request: httpx.Request) -> httpx.Response: assert manager.credential_store.load() is None -def test_device_flow_matches_xai_surface_and_handles_pending(monkeypatch, tmp_path: Path): +def test_device_flow_matches_xai_surface_and_handles_pending( + monkeypatch, tmp_path: Path +): requests: list[httpx.Request] = [] def handler(request: httpx.Request) -> httpx.Response: @@ -791,7 +794,9 @@ def handler(request: httpx.Request) -> httpx.Response: ] results: list[GrokOAuthCredentials] = [] threads = [ - threading.Thread(target=lambda manager=m: results.append(manager.get_valid_credentials())) + threading.Thread( + target=lambda manager=m: results.append(manager.get_valid_credentials()) + ) for m in managers ] for thread in threads: @@ -939,8 +944,7 @@ def handler(request: httpx.Request) -> httpx.Response: 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 + request.headers[GROK_CLI_COMPAT_VERSION_HEADER] == GROK_CLI_COMPAT_VERSION ) return httpx.Response(401 if len(seen) == 1 else 200, json={}) @@ -1158,10 +1162,7 @@ def handler(request: httpx.Request) -> httpx.Response: 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 - ) + 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 @@ -1252,8 +1253,7 @@ async def collect_async(): # type: ignore[no-untyped-def] 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 + request.headers[GROK_CLI_COMPAT_VERSION_HEADER] == GROK_CLI_COMPAT_VERSION ) assert {"temperature", "top_p", "reasoning"}.isdisjoint(payload) @@ -1281,9 +1281,7 @@ def handler(request: httpx.Request) -> httpx.Response: ) async def run() -> str: - async_http_client = httpx.AsyncClient( - transport=httpx.MockTransport(handler) - ) + async_http_client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) llm = GrokOAuth( oauth_credential_path=str(tmp_path / "auth.json"), oauth_access_token="adapter-access", @@ -1308,10 +1306,7 @@ async def run() -> str: 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 - ) + 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): @@ -1351,8 +1346,7 @@ def test_oauth_responses_adapter_pins_proxy_and_omits_controls(tmp_path: Path): 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 + 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 @@ -1437,9 +1431,7 @@ def fake_parse(**kwargs): # type: ignore[no-untyped-def] assert llm.store is False -def test_async_structured_predict_sanitizes_runtime_kwargs( - monkeypatch, tmp_path: Path -): +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", @@ -1495,9 +1487,9 @@ async def fake_parse(**kwargs): # type: ignore[no-untyped-def] 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") + source = Path("mobilerun/agent/utils/oauth/grok_oauth_llm.py").read_text( + encoding="utf-8" + ) forbidden = ( "." + "grok" + "/" + "auth.json", "GROK" + "_HOME", From c76e81fd897251f1d6862cfb889823430a30ba49 Mon Sep 17 00:00:00 2001 From: "rasul.osmanbayli" Date: Thu, 13 Aug 2026 14:08:07 +0400 Subject: [PATCH 4/4] Fix OAuth login timeout handling --- .../agent/utils/oauth/anthropic_oauth_llm.py | 101 +++++-- .../oauth/gemini_oauth_code_assist_llm.py | 137 +++++++-- mobilerun/agent/utils/oauth/grok_oauth_llm.py | 261 ++++++++++++++--- mobilerun/agent/utils/oauth/login_timeout.py | 79 ++++++ .../agent/utils/oauth/openai_oauth_llm.py | 267 +++++++++++++----- mobilerun/cli/main.py | 55 +++- mobilerun/cli/oauth_actions.py | 83 +++++- .../config_manager/auth_profile_store.py | 56 +++- tests/test_grok_cli.py | 101 +++++++ tests/test_grok_oauth.py | 185 ++++++++++-- tests/test_oauth_login_timeout.py | 263 +++++++++++++++++ 11 files changed, 1376 insertions(+), 212 deletions(-) create mode 100644 mobilerun/agent/utils/oauth/login_timeout.py create mode 100644 tests/test_oauth_login_timeout.py diff --git a/mobilerun/agent/utils/oauth/anthropic_oauth_llm.py b/mobilerun/agent/utils/oauth/anthropic_oauth_llm.py index 29de8e06..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,10 @@ 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 @@ -263,7 +268,11 @@ 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() @@ -279,6 +288,8 @@ def _persist_credentials(self) -> 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]: @@ -343,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", @@ -355,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: @@ -371,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: @@ -380,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( @@ -414,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", "" @@ -424,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)) @@ -445,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() @@ -462,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 @@ -479,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( @@ -490,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']}") @@ -511,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, @@ -530,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" @@ -553,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() @@ -576,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.") @@ -608,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 455b0a70..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,10 @@ 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 @@ -292,7 +297,11 @@ 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 @@ -309,6 +318,8 @@ def _persist_credentials(self) -> None: 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]: @@ -328,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 @@ -336,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 [] @@ -364,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: @@ -415,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", @@ -426,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: @@ -441,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: @@ -450,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( @@ -487,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", "" @@ -497,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() @@ -513,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() @@ -530,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 @@ -546,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( @@ -558,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']}") @@ -579,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, @@ -593,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" @@ -612,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() @@ -635,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.") @@ -662,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 index 30aa7ccb..e143bdf3 100644 --- a/mobilerun/agent/utils/oauth/grok_oauth_llm.py +++ b/mobilerun/agent/utils/oauth/grok_oauth_llm.py @@ -18,7 +18,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, Callable, Iterator from urllib.parse import parse_qs, urlencode, urlparse @@ -34,6 +35,10 @@ 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 @@ -239,9 +244,19 @@ def load(self) -> GrokOAuthCredentials | None: return None return GrokOAuthCredentials.from_payload(payload) - def save(self, credentials: GrokOAuthCredentials) -> None: + 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() + 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, ) @@ -250,9 +265,75 @@ class GrokIDTokenValidator: 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() - def validate(self, token: str, *, nonce: str | None) -> dict[str, Any]: - signing_key = self._jwks_client.get_signing_key_from_jwt(token).key + 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, @@ -263,6 +344,8 @@ def validate(self, token: str, *, nonce: str | None) -> dict[str, Any]: ) 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 @@ -332,7 +415,10 @@ def _credentials_from_token_response( 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.") @@ -342,7 +428,21 @@ def _credentials_from_token_response( "xAI authorization response did not contain an ID token." ) if isinstance(id_token, str) and id_token: - self.id_token_validator.validate(id_token, nonce=nonce) + 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: @@ -350,13 +450,16 @@ def _credentials_from_token_response( 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.") - return GrokOAuthCredentials( + 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, @@ -366,26 +469,42 @@ def _post_form( 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=self.request_timeout, + timeout=request_timeout, ) except (httpx.ConnectError, httpx.TimeoutException) as exc: + if deadline is not None: + deadline.check() if attempt < len(backoffs): - self.sleep(backoffs[attempt]) + 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): - self.sleep(backoffs[attempt]) + if deadline is not None: + deadline.sleep(backoffs[attempt]) + else: + self.sleep(backoffs[attempt]) continue return response raise AssertionError("unreachable token retry state") @@ -396,6 +515,7 @@ def _post_token( *, retry_transient: bool = False, refresh_request: bool = False, + deadline: OAuthLoginDeadline | None = None, ) -> dict[str, Any]: response = self._post_form( DEFAULT_GROK_OAUTH_TOKEN_URL, @@ -406,6 +526,7 @@ def _post_token( data=data, context="xAI token request", retry_transient=retry_transient, + deadline=deadline, ) if response.status_code >= 400: try: @@ -419,15 +540,40 @@ def _post_token( raise GrokOAuthError( f"xAI token request failed ({error or response.status_code})." ) - return _safe_json_object(response, context="xAI token response") + 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) -> None: - with self._thread_lock: - self.credential_store.save(credentials) + 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 + self, + *, + code: str, + redirect_uri: str, + code_verifier: str, + nonce: str, + deadline: OAuthLoginDeadline | None = None, ) -> GrokOAuthCredentials: payload = self._post_token( { @@ -436,10 +582,15 @@ def exchange_authorization_code( "code": code, "redirect_uri": redirect_uri, "code_verifier": code_verifier, - } + }, + deadline=deadline, + ) + credentials = self._credentials_from_token_response( + payload, + nonce=nonce, + deadline=deadline, ) - credentials = self._credentials_from_token_response(payload, nonce=nonce) - self.set_initial_credentials(credentials) + self.set_initial_credentials(credentials, deadline=deadline) return credentials def _refresh(self, credentials: GrokOAuthCredentials) -> GrokOAuthCredentials: @@ -757,7 +908,13 @@ def login( 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: @@ -765,7 +922,10 @@ def login( 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(timeout_seconds=timeout_seconds) + 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)) @@ -776,6 +936,7 @@ def login( "error": None, } done = threading.Event() + callback_lock = threading.Lock() class _CallbackHandler(BaseHTTPRequestHandler): def do_GET(self) -> None: # noqa: N802 @@ -784,12 +945,16 @@ def do_GET(self) -> None: # noqa: N802 self.send_response(404) self.end_headers() return - 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"]) + 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() @@ -798,7 +963,6 @@ def do_GET(self) -> None: # noqa: N802 if ok else b"Mobilerun XAI login failed. Return to the terminal." ) - done.set() def log_message(self, format: str, *args: Any) -> None: # noqa: A003 return @@ -806,7 +970,10 @@ def log_message(self, format: str, *args: Any) -> None: # noqa: A003 try: server = HTTPServer((DEFAULT_GROK_OAUTH_CALLBACK_HOST, 0), _CallbackHandler) except OSError: - return self._login_device_code(timeout_seconds=timeout_seconds) + 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]}" @@ -819,13 +986,15 @@ def log_message(self, format: str, *args: Any) -> None: # noqa: A003 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: - webbrowser.open(authorization_url) - if not done.wait(timeout=max(0.0, timeout_seconds)): + 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'})." @@ -839,16 +1008,20 @@ def log_message(self, format: str, *args: Any) -> None: # noqa: A003 redirect_uri=redirect_uri, code_verifier=code_verifier, nonce=nonce, + deadline=login_deadline, ) finally: server.shutdown() server.server_close() def _login_device_code( - self, *, timeout_seconds: float = 1800.0 + self, + *, + deadline: OAuthLoginDeadline, + open_browser: bool, ) -> GrokOAuthCredentials: manager = self._oauth_manager - response = manager.http_client.post( + response = manager._post_form( DEFAULT_GROK_OAUTH_DEVICE_URL, headers={ "Accept": "application/json", @@ -859,7 +1032,9 @@ def _login_device_code( "client_id": DEFAULT_GROK_OAUTH_CLIENT_ID, "scope": " ".join(DEFAULT_GROK_OAUTH_SCOPES), }, - timeout=manager.request_timeout, + context="xAI device authorization request", + retry_transient=False, + deadline=deadline, ) if response.status_code >= 400: raise GrokOAuthError( @@ -883,14 +1058,17 @@ def _login_device_code( interval = max(1, int(payload.get("interval", 5))) except (TypeError, ValueError): expires_in, interval = 1800, 5 - deadline = time.monotonic() + min(max(0.0, timeout_seconds), expires_in) + 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 time.monotonic() < deadline: + while True: + device_deadline.check() token_response = manager._post_form( DEFAULT_GROK_OAUTH_TOKEN_URL, headers={ @@ -904,13 +1082,21 @@ def _login_device_code( }, 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" ) - credentials = manager._credentials_from_token_response(token_payload) - manager.set_initial_credentials(credentials) + 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")) @@ -928,8 +1114,7 @@ def _login_device_code( raise GrokOAuthError( f"xAI device token request failed ({error or token_response.status_code})." ) - manager.sleep(min(interval, max(0.0, deadline - time.monotonic()))) - raise TimeoutError("XAI OAuth device authorization timed out.") + device_deadline.sleep(interval) # Descriptive alias for callers that prefer the full class name. 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 1c89b8b1..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,10 @@ 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 @@ -57,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: @@ -81,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). @@ -89,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", @@ -117,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}") @@ -135,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.") @@ -169,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 @@ -290,8 +332,18 @@ def load(self) -> Optional[OpenAIOAuthCredentials]: except ValueError: return None - def save(self, credentials: OpenAIOAuthCredentials) -> None: - self._store.update_slot(self._NESTED_KEY, credentials.to_dict()) + 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: @@ -365,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: @@ -449,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", @@ -458,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: @@ -493,7 +572,7 @@ def exchange_authorization_code( else access ), ) - self.set_initial_credentials(credentials) + self.set_initial_credentials(credentials, deadline=deadline) return credentials @@ -641,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)) @@ -665,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() @@ -682,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 @@ -690,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}" @@ -708,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']}") @@ -733,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) @@ -744,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. @@ -752,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: @@ -767,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") @@ -793,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, @@ -801,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") @@ -817,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/main.py b/mobilerun/cli/main.py index c76d101d..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 @@ -77,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 @@ -418,9 +429,11 @@ 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) @@ -1088,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", @@ -1150,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") @@ -1172,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", @@ -1235,9 +1273,10 @@ def configure_gemini( @click.option( "--timeout", type=float, + callback=_validate_oauth_timeout, default=300.0, show_default=True, - help="Max seconds to wait for xAI OAuth authentication.", + help="Max seconds allowed for the complete xAI OAuth login.", ) @click.option( "--open-browser/--no-browser", diff --git a/mobilerun/cli/oauth_actions.py b/mobilerun/cli/oauth_actions.py index 2e29ad59..2ffd1144 100644 --- a/mobilerun/cli/oauth_actions.py +++ b/mobilerun/cli/oauth_actions.py @@ -14,6 +14,7 @@ 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, @@ -35,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, ) @@ -55,28 +65,45 @@ 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.") @@ -86,44 +113,71 @@ def run_grok_oauth_login( timeout: float = 300.0, open_browser: bool = True, device_code: bool = False, - no_browser: 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 and not no_browser, - timeout_seconds=timeout, + 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: +def save_anthropic_setup_token( + credential_path: str, + token: str, + *, + deadline: OAuthLoginDeadline | None = None, +) -> None: AuthProfileStore(credential_path).update_slot( "claudeAiOauth", { @@ -132,12 +186,13 @@ def save_anthropic_setup_token(credential_path: str, token: str) -> 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 index e436091e..5b50d158 100644 --- a/mobilerun/config_manager/auth_profile_store.py +++ b/mobilerun/config_manager/auth_profile_store.py @@ -28,15 +28,26 @@ class AuthProfileFormatError(ValueError): class AuthProfileTransaction(AbstractContextManager["AuthProfileTransaction"]): """A locked read/modify/write transaction over an auth profile file.""" - def __init__(self, store: "AuthProfileStore") -> None: + 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) - self._lock.acquire() + 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. @@ -71,7 +82,10 @@ def update(self, updater: Callable[[dict[str, Any]], _T]) -> _T: 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) + self._store._write_unlocked( + self._profile, + before_commit=self._before_commit, + ) finally: self._lock.release() return None @@ -84,8 +98,17 @@ 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) -> AuthProfileTransaction: - return AuthProfileTransaction(self) + 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: @@ -95,8 +118,18 @@ 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]) -> None: - with self.transaction() as transaction: + 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: @@ -119,7 +152,12 @@ def _read_unlocked(self) -> dict[str, Any]: ) return payload - def _write_unlocked(self, profile: dict[str, Any]) -> None: + 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 @@ -138,6 +176,8 @@ def _write_unlocked(self, profile: dict[str, Any]) -> None: 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) diff --git a/tests/test_grok_cli.py b/tests/test_grok_cli.py index 8ce8ab76..7e14ddbf 100644 --- a/tests/test_grok_cli.py +++ b/tests/test_grok_cli.py @@ -8,6 +8,8 @@ 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 @@ -100,6 +102,105 @@ def test_configure_xai_command_forwards_device_code_options( ] +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() diff --git a/tests/test_grok_oauth.py b/tests/test_grok_oauth.py index 8f9edcdc..b7fb4d7c 100644 --- a/tests/test_grok_oauth.py +++ b/tests/test_grok_oauth.py @@ -43,6 +43,7 @@ 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, @@ -59,6 +60,26 @@ def validate(self, token: str, *, nonce: str | None): # type: ignore[no-untyped 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", @@ -289,16 +310,74 @@ def test_id_token_validator_checks_es256_claims_and_nonce(): algorithm="ES256", headers={"kid": "test"}, ) - jwks_client = SimpleNamespace( - get_signing_key_from_jwt=lambda _: SimpleNamespace(key=public_key) - ) + 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) - assert validator.validate(token, nonce="expected")["sub"] == "user" + 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()) @@ -409,6 +488,8 @@ def test_callback_query_rejects_duplicate_code_state_and_error_values(): 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: @@ -468,6 +549,7 @@ def server_close(self) -> None: 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:") @@ -491,7 +573,11 @@ def make_server(address, handler): # type: ignore[no-untyped-def] "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) + credentials = llm.login( + open_browser=True, + timeout_seconds=5, + deadline=deadline, + ) assert credentials.access_token == "access" assert len(opened) == 1 @@ -499,6 +585,7 @@ def make_server(address, handler): # type: ignore[no-untyped-def] 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): @@ -632,16 +719,20 @@ def test_token_errors_never_expose_response_bodies( 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, ) @@ -652,18 +743,23 @@ def handler(request: httpx.Request) -> httpx.Response: 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, tmp_path: Path + 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) @@ -673,7 +769,7 @@ def handler(request: httpx.Request) -> httpx.Response: json={ "device_code": "device-secret", "user_code": "ABCD-1234", - "verification_uri_complete": "https://accounts.x.ai/device?code=ABCD-1234", + "verification_uri_complete": verification_uri, "expires_in": 1800, "interval": 1, }, @@ -690,6 +786,10 @@ def handler(request: httpx.Request) -> httpx.Response: ) 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)), @@ -697,7 +797,11 @@ def handler(request: httpx.Request) -> httpx.Response: ) llm = GrokOAuth(oauth_session_manager=manager) - credentials = llm.login(device_code=True, timeout_seconds=10) + 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" @@ -711,15 +815,51 @@ def handler(request: httpx.Request) -> httpx.Response: 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_device_token_poll_retries_connect_and_5xx_failures(tmp_path: Path): +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 - delays: list[float] = [] + 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={ @@ -732,9 +872,12 @@ def handler(request: httpx.Request) -> httpx.Response: ) 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={ @@ -744,22 +887,26 @@ def handler(request: httpx.Request) -> httpx.Response: }, ) + store = GrokOAuthCredentialStore(tmp_path / "auth.json") manager = GrokOAuthSessionManager( - credential_store=GrokOAuthCredentialStore(tmp_path / "auth.json"), + 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, + request_timeout=20, + retry_backoff_seconds=(0.25, 0.5), ) - credentials = GrokOAuth(oauth_session_manager=manager).login( - device_code=True, - timeout_seconds=10, - ) + 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 delays == [0.01, 0.02] - assert credentials.access_token == "access" + 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( 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