From a460a2f50a36ec2b021507edd8e4c761628021ec Mon Sep 17 00:00:00 2001
From: srojk34 <286497132+srojk34@users.noreply.github.com>
Date: Wed, 1 Jul 2026 00:17:59 +0300
Subject: [PATCH] fix(xai): forward image_gen.model kwarg to _resolve_model in
 generate()

generate() accepted **kwargs but called _resolve_model() with no
arguments, silently discarding any model the caller supplied.
tools/image_generation_tool.py passes kwargs["model"] = configured_model
sourced from image_gen.model in config.yaml into provider.generate().
For the xAI provider this meant a user selecting grok-imagine-image-quality
via hermes tools always got the default grok-imagine-image instead, with
no error or log. Mirrors the fix PR #55672 applied to the openrouter
provider, which already threads kwargs.get("model") into its resolver.

Add caller_model parameter to _resolve_model() at the highest priority
(before env override), and pass kwargs.get("model") from generate().
An unrecognised model name falls through to the existing priority chain.
---
 plugins/image_gen/xai/__init__.py            | 20 +++++++++--
 tests/plugins/image_gen/test_xai_provider.py | 37 ++++++++++++++++++++
 2 files changed, 54 insertions(+), 3 deletions(-)

diff --git a/plugins/image_gen/xai/__init__.py b/plugins/image_gen/xai/__init__.py
index 31a0b719bb6f..cff2034e613b 100644
--- a/plugins/image_gen/xai/__init__.py
+++ b/plugins/image_gen/xai/__init__.py
@@ -100,8 +100,22 @@ def _load_xai_config() -> Dict[str, Any]:
         return {}
 
 
-def _resolve_model() -> Tuple[str, Dict[str, Any]]:
-    """Decide which model to use and return ``(model_id, meta)``."""
+def _resolve_model(caller_model: Optional[str] = None) -> Tuple[str, Dict[str, Any]]:
+    """Decide which model to use and return ``(model_id, meta)``.
+
+    Priority:
+    1. Caller-supplied ``caller_model`` (from ``generate(**kwargs)`` via
+       ``image_gen.model`` config key) — mirrors how the openrouter provider
+       threads ``kwargs.get("model")`` into its own resolver.
+    2. ``XAI_IMAGE_MODEL`` env override.
+    3. ``image_gen.model`` in config.yaml (already read by the caller into
+       ``caller_model``; this branch handles the legacy path where config is
+       read directly here rather than forwarded by the caller).
+    4. Hard-coded default.
+    """
+    if caller_model and caller_model in _MODELS:
+        return caller_model, _MODELS[caller_model]
+
     env_override = os.environ.get("XAI_IMAGE_MODEL")
     if env_override and env_override in _MODELS:
         return env_override, _MODELS[env_override]
@@ -234,7 +248,7 @@ def generate(
                 aspect_ratio=aspect_ratio,
             )
 
-        model_id, meta = _resolve_model()
+        model_id, meta = _resolve_model(kwargs.get("model"))
         aspect = resolve_aspect_ratio(aspect_ratio)
         xai_ar = _XAI_ASPECT_RATIOS.get(aspect, "1:1")
         resolution = _resolve_resolution()
diff --git a/tests/plugins/image_gen/test_xai_provider.py b/tests/plugins/image_gen/test_xai_provider.py
index cf9708dae1d1..a7833c3490c5 100644
--- a/tests/plugins/image_gen/test_xai_provider.py
+++ b/tests/plugins/image_gen/test_xai_provider.py
@@ -120,6 +120,43 @@ def test_custom_model(self, monkeypatch):
         model_id, _ = _resolve_model()
         assert model_id == "grok-imagine-image"
 
+    def test_caller_model_overrides_env(self, monkeypatch):
+        """caller_model (from image_gen.model config key) must take priority
+        over XAI_IMAGE_MODEL env — mirrors the fix applied to the openrouter
+        provider in #55672."""
+        monkeypatch.setenv("XAI_IMAGE_MODEL", "grok-imagine-image")
+        from plugins.image_gen.xai import _resolve_model
+
+        model_id, _ = _resolve_model("grok-imagine-image-quality")
+        assert model_id == "grok-imagine-image-quality"
+
+    def test_unknown_caller_model_falls_back_to_env(self, monkeypatch):
+        """An unrecognised caller_model must not crash — fall through to env."""
+        monkeypatch.setenv("XAI_IMAGE_MODEL", "grok-imagine-image")
+        from plugins.image_gen.xai import _resolve_model
+
+        model_id, _ = _resolve_model("not-a-real-model")
+        assert model_id == "grok-imagine-image"
+
+    def test_model_kwarg_forwarded_to_generate(self):
+        """generate(model=...) must use the supplied model, not the default."""
+        from plugins.image_gen.xai import XAIImageGenProvider
+
+        mock_resp = MagicMock()
+        mock_resp.status_code = 200
+        mock_resp.raise_for_status = MagicMock()
+        mock_resp.json.return_value = {"data": [{"b64_json": "dGVzdA=="}]}
+
+        with patch("plugins.image_gen.xai.requests.post", return_value=mock_resp) as mock_post:
+            with patch("plugins.image_gen.xai.save_b64_image", return_value="/tmp/out.png"):
+                provider = XAIImageGenProvider()
+                result = provider.generate(prompt="test", model="grok-imagine-image-quality")
+
+        assert result["success"] is True
+        assert result["model"] == "grok-imagine-image-quality"
+        payload = mock_post.call_args.kwargs.get("json") or mock_post.call_args[1].get("json", {})
+        assert payload.get("model") == "grok-imagine-image-quality"
+
 
 # ---------------------------------------------------------------------------
 # Generate tests