Studio: fix backend CI test failures

- test_safetensors_tool_loop: FakeExecuteTool stub now accepts the
  tool_context kwarg the loop passes (production is correct).
- test_desktop_auth: add rag_router to the health-check router stub
  so main.py's router import resolves.
- test_rag_reingest / test_rag_multimodal: load routes/rag.py by file
  path instead of `from routes.rag import`, which runs routes/__init__
  and eagerly imports the datasets router. On the GPU-less repo-cpu
  runner the unsloth bootstrap can leave `datasets` half-initialized,
  making that eager `from datasets import IterableDataset` raise.

studio/backend/tests/test_desktop_auth.py

@@ -265,9 +265,9 @@ def test_consume_refresh_token_concurrent_only_one_succeeds(tmp_path, monkeypatc
         results = list(pool.map(attempt, range(workers)))
 
     successes = [r for r in results if r is not None]
-    assert (
-        len(successes) == 1
-    ), f"expected exactly one consumer to win, got {len(successes)}"
+    assert len(successes) == 1, (
+        f"expected exactly one consumer to win, got {len(successes)}"
+    )
     assert successes[0] == (storage.DEFAULT_ADMIN_USERNAME, False)
 
 
@@ -440,6 +440,7 @@ def test_health_response_reports_desktop_capability_fields(monkeypatch):
         mcp_servers_router = APIRouter(),
         models_router = APIRouter(),
         providers_router = APIRouter(),
+        rag_router = APIRouter(),
         training_history_router = APIRouter(),
         training_router = APIRouter(),
     )

studio/backend/tests/test_safetensors_tool_loop.py

@@ -172,6 +172,7 @@ class FakeExecuteTool:
         cancel_event = None,
         timeout = None,
         session_id = None,
+        tool_context = None,
     ):
         self.calls.append((name, arguments))
         result = self.results.pop(0) if self.results else "OK"
@@ -384,9 +385,9 @@ class TestLoopBehaviour:
         tool_msgs = [m for m in captured[1] if m.get("role") == "tool"]
         assert tool_msgs, "no tool message reached the model"
         for tm in tool_msgs:
-            assert (
-                "__IMAGES__" not in tm["content"]
-            ), f"sentinel leaked to model: {tm['content']!r}"
+            assert "__IMAGES__" not in tm["content"], (
+                f"sentinel leaked to model: {tm['content']!r}"
+            )
 
     def test_image_sentinel_stripped_with_multiple_markers(self):
         # Consecutive sentinels: cut at the first, nothing leaks.
@@ -416,12 +417,12 @@ class TestLoopBehaviour:
         tool_msgs = [m for m in captured[1] if m.get("role") == "tool"]
         assert tool_msgs
         for tm in tool_msgs:
-            assert (
-                "__IMAGES__" not in tm["content"]
-            ), f"second sentinel leaked: {tm['content']!r}"
-            assert (
-                tm["content"] == "panel"
-            ), f"expected payload-only 'panel', got {tm['content']!r}"
+            assert "__IMAGES__" not in tm["content"], (
+                f"second sentinel leaked: {tm['content']!r}"
+            )
+            assert tm["content"] == "panel", (
+                f"expected payload-only 'panel', got {tm['content']!r}"
+            )
 
     def test_tool_execution_error_is_emitted_but_loop_continues(self):
         loop, exec_fn = _make_loop(
@@ -544,9 +545,9 @@ class TestProseMentioningToolCall:
         contents = [e for e in events if e["type"] == "content"]
         assert contents, "expected at least one content event"
         final = contents[-1]["text"]
-        assert (
-            "LLM tool" in final
-        ), f"prose mentioning <tool_call> should not be truncated; got {final!r}"
+        assert "LLM tool" in final, (
+            f"prose mentioning <tool_call> should not be truncated; got {final!r}"
+        )
 
     def test_tool_result_with_tool_call_text_does_not_retrigger(self):
         # Tool result text contains the literal ``<tool_call>`` string.

tests/python/test_rag_multimodal.py

@@ -7,6 +7,7 @@ returns images when asked, route accepts the mode field, constraint
 validator rejects illegal combos) run in every test invocation.
 """
 
+import importlib.util
 import sys
 from pathlib import Path
 
@@ -18,6 +19,27 @@ if str(STUDIO_BACKEND) not in sys.path:
     sys.path.insert(0, str(STUDIO_BACKEND))
 
 
+def _rag_route():
+    """Load ``routes/rag.py`` directly, bypassing the ``routes`` package.
+
+    ``from routes.rag import X`` first runs ``routes/__init__.py``, which eagerly
+    imports every router — including the datasets router, whose chain does
+    ``from datasets import IterableDataset`` at import time. On a GPU-less CI
+    runner the unsloth bootstrap can leave ``datasets`` half-initialized, so that
+    eager import raises. These tests only need pure helpers from rag.py, so load
+    the file on its own (it has no intra-``routes`` imports).
+    """
+    mod = sys.modules.get("_rag_route_under_test")
+    if mod is None:
+        spec = importlib.util.spec_from_file_location(
+            "_rag_route_under_test", STUDIO_BACKEND / "routes" / "rag.py"
+        )
+        mod = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(mod)
+        sys.modules["_rag_route_under_test"] = mod
+    return mod
+
+
 def test_html_parser_returns_images_when_requested(tmp_path):
     pytest.importorskip("bs4")
     pytest.importorskip("lxml")
@@ -55,7 +77,7 @@ def test_html_parser_returns_images_when_requested(tmp_path):
 def test_multimodal_late_combo_validator():
     from fastapi import HTTPException
 
-    from routes.rag import _validate_mode_combo
+    _validate_mode_combo = _rag_route()._validate_mode_combo
 
     # Allowed combos → None.
     assert _validate_mode_combo("text", "standard") is None

tests/python/test_rag_reingest.py

@@ -6,6 +6,7 @@ parts that are testable without external models: payload validation and
 the (multimodal, late) constraint propagation.
 """
 
+import importlib.util
 import sys
 from pathlib import Path
 
@@ -17,8 +18,29 @@ if str(STUDIO_BACKEND) not in sys.path:
     sys.path.insert(0, str(STUDIO_BACKEND))
 
 
+def _rag_route():
+    """Load ``routes/rag.py`` directly, bypassing the ``routes`` package.
+
+    ``from routes.rag import X`` first runs ``routes/__init__.py``, which eagerly
+    imports every router — including the datasets router, whose chain does
+    ``from datasets import IterableDataset`` at import time. On a GPU-less CI
+    runner the unsloth bootstrap can leave ``datasets`` half-initialized, so that
+    eager import raises. These tests only need pure helpers from rag.py, so load
+    the file on its own (it has no intra-``routes`` imports).
+    """
+    mod = sys.modules.get("_rag_route_under_test")
+    if mod is None:
+        spec = importlib.util.spec_from_file_location(
+            "_rag_route_under_test", STUDIO_BACKEND / "routes" / "rag.py"
+        )
+        mod = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(mod)
+        sys.modules["_rag_route_under_test"] = mod
+    return mod
+
+
 def test_reingest_request_accepts_all_optional_fields():
-    from routes.rag import ReingestKBRequest
+    ReingestKBRequest = _rag_route().ReingestKBRequest
 
     empty = ReingestKBRequest()
     assert empty.chunking_strategy is None
@@ -31,17 +53,19 @@ def test_reingest_request_accepts_all_optional_fields():
 
 
 def test_reingest_request_rejects_unknown_strategy():
-    from routes.rag import ReingestKBRequest
     from pydantic import ValidationError
 
+    ReingestKBRequest = _rag_route().ReingestKBRequest
+
     with pytest.raises(ValidationError):
         ReingestKBRequest(chunking_strategy = "telekinetic")
 
 
 def test_reingest_request_rejects_unknown_mode():
-    from routes.rag import ReingestKBRequest
     from pydantic import ValidationError
 
+    ReingestKBRequest = _rag_route().ReingestKBRequest
+
     with pytest.raises(ValidationError):
         ReingestKBRequest(mode = "augmented")
 
@@ -50,7 +74,7 @@ def test_constraint_still_enforced_for_reingest_combos():
     """The combination guard is shared with create — verify it still bites."""
     from fastapi import HTTPException
 
-    from routes.rag import _validate_mode_combo
+    _validate_mode_combo = _rag_route()._validate_mode_combo
 
     with pytest.raises(HTTPException) as excinfo:
         _validate_mode_combo("multimodal", "late")