From 09afafb87e486eea62de0d7892a127808f121b0a Mon Sep 17 00:00:00 2001
From: Julien Talbot <61507425+Julientalbot@users.noreply.github.com>
Date: Fri, 22 May 2026 21:57:43 +0400
Subject: [PATCH] fix(xai): resolve Grok Build context for OAuth

---
 agent/model_metadata.py            |  1 +
 agent/models_dev.py                |  3 +++
 tests/agent/test_model_metadata.py | 28 ++++++++++++++++++++++++++++
 tests/agent/test_models_dev.py     | 20 ++++++++++++++++++++
 4 files changed, 52 insertions(+)

diff --git a/agent/model_metadata.py b/agent/model_metadata.py
index b8ec0d6509e..3d6216f6beb 100644
--- a/agent/model_metadata.py
+++ b/agent/model_metadata.py
@@ -209,6 +209,7 @@ DEFAULT_CONTEXT_LENGTHS = {
     # via a custom provider. Values sourced from models.dev (2026-04).
     # Keys use substring matching (longest-first), so e.g. "grok-4.20"
     # matches "grok-4.20-0309-reasoning" / "-non-reasoning" / "-multi-agent-0309".
+    "grok-build": 256000,       # grok-build-0.1
     "grok-code-fast": 256000,   # grok-code-fast-1
     "grok-4-1-fast": 2000000,   # grok-4-1-fast-(non-)reasoning
     "grok-2-vision": 8192,      # grok-2-vision, -1212, -latest
diff --git a/agent/models_dev.py b/agent/models_dev.py
index 8fabb276645..1249c6f1970 100644
--- a/agent/models_dev.py
+++ b/agent/models_dev.py
@@ -167,6 +167,9 @@ PROVIDER_TO_MODELS_DEV: Dict[str, str] = {
     "gemini": "google",
     "google": "google",
     "xai": "xai",
+    # xAI OAuth is an authentication/transport path for the same xAI model
+    # catalog, so model metadata should resolve through the xAI provider.
+    "xai-oauth": "xai",
     "xiaomi": "xiaomi",
     "nvidia": "nvidia",
     "groq": "groq",
diff --git a/tests/agent/test_model_metadata.py b/tests/agent/test_model_metadata.py
index 4f2b51293a6..e905c3e1f6b 100644
--- a/tests/agent/test_model_metadata.py
+++ b/tests/agent/test_model_metadata.py
@@ -164,6 +164,7 @@ class TestDefaultContextLengths:
             "grok-4-1-fast": 2000000,
             "grok-4-fast": 2000000,
             "grok-4": 256000,
+            "grok-build": 256000,
             "grok-code-fast": 256000,
             "grok-3": 131072,
             "grok-2": 131072,
@@ -195,6 +196,7 @@ class TestDefaultContextLengths:
                 ("grok-4-fast-non-reasoning", 2000000),
                 ("grok-4", 256000),
                 ("grok-4-0709", 256000),
+                ("grok-build-0.1", 256000),
                 ("grok-code-fast-1", 256000),
                 ("grok-3", 131072),
                 ("grok-3-mini", 131072),
@@ -210,6 +212,32 @@ class TestDefaultContextLengths:
                     f"{model_id}: expected {expected_ctx}, got {actual}"
                 )
 
+    def test_xai_oauth_grok_build_uses_xai_models_dev_context(self):
+        """xAI OAuth should share the xAI provider metadata path.
+
+        The xAI /v1/models endpoint does not currently include context fields
+        for grok-build-0.1, so this guards against falling through to the
+        generic "grok" 131k fallback when using OAuth credentials.
+        """
+        registry = {
+            "xai": {
+                "models": {
+                    "grok-build-0.1": {
+                        "limit": {"context": 256000, "output": 64000},
+                    },
+                },
+            },
+        }
+        with patch("agent.model_metadata.get_cached_context_length", return_value=None), \
+             patch("agent.model_metadata._query_ollama_api_show", return_value=None), \
+             patch("agent.models_dev.fetch_models_dev", return_value=registry):
+            assert get_model_context_length(
+                "grok-build-0.1",
+                provider="xai-oauth",
+                base_url="https://api.x.ai/v1",
+                api_key="oauth-token",
+            ) == 256000
+
     def test_deepseek_v4_models_1m_context(self):
         from agent.model_metadata import get_model_context_length
         from unittest.mock import patch as mock_patch
diff --git a/tests/agent/test_models_dev.py b/tests/agent/test_models_dev.py
index 2cb9746b223..e3338091b9f 100644
--- a/tests/agent/test_models_dev.py
+++ b/tests/agent/test_models_dev.py
@@ -41,6 +41,16 @@ SAMPLE_REGISTRY = {
             },
         },
     },
+    "xai": {
+        "id": "xai",
+        "name": "xAI",
+        "models": {
+            "grok-build-0.1": {
+                "id": "grok-build-0.1",
+                "limit": {"context": 256000, "output": 64000},
+            },
+        },
+    },
     "kilo": {
         "id": "kilo",
         "name": "Kilo Gateway",
@@ -86,6 +96,10 @@ class TestProviderMapping:
         assert PROVIDER_TO_MODELS_DEV["kilocode"] == "kilo"
         assert PROVIDER_TO_MODELS_DEV["ai-gateway"] == "vercel"
 
+    def test_xai_oauth_uses_xai_catalog(self):
+        assert PROVIDER_TO_MODELS_DEV["xai"] == "xai"
+        assert PROVIDER_TO_MODELS_DEV["xai-oauth"] == "xai"
+
     def test_unmapped_provider_not_in_dict(self):
         assert "nous" not in PROVIDER_TO_MODELS_DEV
 
@@ -144,6 +158,12 @@ class TestLookupModelsDevContext:
         # GitHub Copilot: only 128K for same model
         assert lookup_models_dev_context("copilot", "claude-opus-4.6") == 128000
 
+    @patch("agent.models_dev.fetch_models_dev")
+    def test_xai_oauth_resolves_xai_context(self, mock_fetch):
+        """xAI OAuth is an auth path, not a separate model catalog."""
+        mock_fetch.return_value = SAMPLE_REGISTRY
+        assert lookup_models_dev_context("xai-oauth", "grok-build-0.1") == 256000
+
     @patch("agent.models_dev.fetch_models_dev")
     def test_zero_context_filtered(self, mock_fetch):
         mock_fetch.return_value = SAMPLE_REGISTRY