fix(server): correct thread and model injection errors (#2240)

Fixes two server issues:

Return HTTP 400 when thread_id is used without a configured datastore.
Preserve the main model’s api_key_env_var during request model injection.

Signed-off-by: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com>

Author: Pouyanpi

nemoguardrails/server/api.py

@@ -381,9 +381,14 @@ def _update_models_in_config(config: RailsConfig, main_model: Model) -> RailsCon
             break
 
     if main_model_index is not None:
-        parameters = {**models[main_model_index].parameters, **main_model.parameters}
-        models[main_model_index] = main_model
-        models[main_model_index].parameters = parameters
+        configured_model = models[main_model_index]
+        parameters = {**configured_model.parameters, **main_model.parameters}
+        models[main_model_index] = main_model.model_copy(
+            update={
+                "api_key_env_var": main_model.api_key_env_var or configured_model.api_key_env_var,
+                "parameters": parameters,
+            }
+        )
     else:
         models.append(main_model)
 
@@ -636,12 +641,9 @@ async def chat_completion(body: GuardrailsChatCompletionRequest, request: Reques
 
     if body.guardrails.thread_id:
         if datastore is None:
-            raise RuntimeError("No DataStore has been configured.")
-        # We make sure the `thread_id` meets the minimum complexity requirement.
-        if len(body.guardrails.thread_id) < 16:
             raise HTTPException(
-                status_code=422,
-                detail="The `thread_id` must have a minimum length of 16 characters.",
+                status_code=400,
+                detail="Conversation threads are not enabled on this server.",
             )
 
         # Fetch the existing thread messages. For easier management, we prepend

tests/server/test_api.py

@@ -24,6 +24,10 @@
 pytest.importorskip("openai", reason="openai is required for server tests")
 from fastapi.testclient import TestClient
 
+from nemoguardrails import RailsConfig
+from nemoguardrails.guardrails.model_engine import ModelEngine
+from nemoguardrails.llm.models.openai_chat import OpenAIChatModel
+from nemoguardrails.rails import LLMRails
 from nemoguardrails.server import api
 from nemoguardrails.server.api import _format_streaming_response
 from nemoguardrails.server.schemas.openai import GuardrailsChatCompletionRequest
@@ -166,6 +170,82 @@ def test_model_field_independent_of_config_id():
     assert request_body.guardrails.config_ids == ["test_config"]
 
 
+@pytest.fixture
+def injected_model_config(monkeypatch):
+    monkeypatch.setenv("CUSTOM_MAIN_API_KEY", "main-key")
+    monkeypatch.setenv("MAIN_MODEL_BASE_URL", "https://request.example/v1")
+    config = RailsConfig.from_content(
+        config={
+            "models": [
+                {
+                    "type": "main",
+                    "engine": "nim",
+                    "model": "configured-model",
+                    "api_key_env_var": "CUSTOM_MAIN_API_KEY",
+                    "parameters": {
+                        "base_url": "https://configured.example/v1",
+                        "default_headers": {"X-Tenant": "acme"},
+                    },
+                }
+            ]
+        }
+    )
+    return api._inject_model(config, "requested-model")
+
+
+def test_inject_model_preserves_main_model_api_key_env_var(injected_model_config):
+    main_model = injected_model_config.models[0]
+    headers = ModelEngine(main_model)._prepare_request([{"role": "user", "content": "hi"}]).headers
+
+    assert main_model.model == "requested-model"
+    assert main_model.engine == "custom_llm"
+    assert main_model.api_key_env_var == "CUSTOM_MAIN_API_KEY"
+    assert main_model.parameters == {
+        "base_url": "https://request.example/v1",
+        "default_headers": {"X-Tenant": "acme"},
+    }
+    assert headers["Authorization"] == "Bearer main-key"
+    assert headers["X-Tenant"] == "acme"
+
+
+def test_inject_model_preserves_main_model_api_key_for_llmrails(injected_model_config):
+    rails = LLMRails(config=injected_model_config.model_copy(deep=True))
+
+    assert isinstance(rails.llm, OpenAIChatModel)
+    headers = rails.llm._client._build_headers()
+    assert rails.llm.model_name == "requested-model"
+    assert rails.llm.provider_name == "custom_llm"
+    assert headers["Authorization"] == "Bearer main-key"
+    assert headers["X-Tenant"] == "acme"
+
+
+def test_thread_id_without_datastore_returns_400(monkeypatch):
+    mock_rails = AsyncMock()
+    mock_rails.config = RailsConfig.from_content(config={"models": []})
+    monkeypatch.setattr(api, "datastore", None)
+
+    with patch("nemoguardrails.server.api._get_rails", new=AsyncMock(return_value=mock_rails)):
+        response = client.post(
+            "/v1/chat/completions",
+            json={
+                "model": "gpt-4o",
+                "messages": [{"role": "user", "content": "Hello"}],
+                "guardrails": {
+                    "config_id": "test_config",
+                    "thread_id": "0123456789abcdef",
+                },
+            },
+        )
+
+    assert response.status_code == 400
+    assert response.json()["error"] == {
+        "message": "Conversation threads are not enabled on this server.",
+        "type": "invalid_request_error",
+        "param": None,
+        "code": None,
+    }
+
+
 def test_request_body_state():
     """Test GuardrailsChatCompletionRequest state handling."""
     data = {