feat: support per-request model override in /v1/chat/completions (#103)

* feat: support per-request model override for /v1/chat/completions - add optional model override to completion request types\n- forward request model through gateway, worker, and orchestrator proxy paths\n- use request model in NEAR AI providers with fallback to active model\n- replace model-mismatch integration test with override propagation checks\n- update FEATURE_PARITY.md note for OpenAI-compatible API behavior\n\nRefs #49 * Wire gateway OpenAI-compatible routes to active LLM provider * Validate OpenAI model name length before streaming * Address PR103 review feedback on model override and validation * Report effective model in OpenAI-compatible responses * Use async mutexes in OpenAI compatibility integration tests * fix tests for per-request model field in response cache * fix formatting and clippy lint after main merge * Fix model override reporting and cache correctness --------- Co-authored-by: Illia Polosukhin <[email protected]>
2026-08-25 14:53:34 +00:00 · 2026-02-19 21:45:37 +00:00
parent 89fdd81420
commit ccf60055f4
13 changed files with 519 additions and 62 deletions
@@ -24,7 +24,21 @@ const AUTH_TOKEN: &str = "test-openai-token";
 // Mock LLM provider
 // ---------------------------------------------------------------------------

-struct MockLlmProvider;
+#[derive(Default)]
+struct MockLlmState {
+    completion_models: tokio::sync::Mutex<Vec<Option<String>>>,
+    tool_completion_models: tokio::sync::Mutex<Vec<Option<String>>>,
+}
+
+struct MockLlmProvider {
+    state: Arc<MockLlmState>,
+}
+
+impl MockLlmProvider {
+    fn new(state: Arc<MockLlmState>) -> Self {
+        Self { state }
+    }
+}

 #[async_trait]
 impl LlmProvider for MockLlmProvider {
@@ -37,6 +51,12 @@ impl LlmProvider for MockLlmProvider {
    }

    async fn complete(&self, req: CompletionRequest) -> Result<CompletionResponse, LlmError> {
+        self.state
+            .completion_models
+            .lock()
+            .await
+            .push(req.model.clone());
+
        // Echo the last user message back
        let user_msg = req
            .messages
@@ -59,6 +79,12 @@ impl LlmProvider for MockLlmProvider {
        &self,
        req: ToolCompletionRequest,
    ) -> Result<ToolCompletionResponse, LlmError> {
+        self.state
+            .tool_completion_models
+            .lock()
+            .await
+            .push(req.model.clone());
+
        // If tools are provided, return a tool call
        if let Some(tool) = req.tools.first() {
            Ok(ToolCompletionResponse {
@@ -93,11 +119,71 @@ impl LlmProvider for MockLlmProvider {
    }
 }

+struct FixedModelProvider {
+    model: &'static str,
+}
+
+impl FixedModelProvider {
+    fn new(model: &'static str) -> Self {
+        Self { model }
+    }
+}
+
+#[async_trait]
+impl LlmProvider for FixedModelProvider {
+    fn model_name(&self) -> &str {
+        self.model
+    }
+
+    fn cost_per_token(&self) -> (Decimal, Decimal) {
+        (Decimal::ZERO, Decimal::ZERO)
+    }
+
+    async fn complete(&self, _req: CompletionRequest) -> Result<CompletionResponse, LlmError> {
+        Ok(CompletionResponse {
+            content: "fixed response".to_string(),
+            input_tokens: 10,
+            output_tokens: 5,
+            finish_reason: FinishReason::Stop,
+            response_id: None,
+        })
+    }
+
+    async fn complete_with_tools(
+        &self,
+        _req: ToolCompletionRequest,
+    ) -> Result<ToolCompletionResponse, LlmError> {
+        Ok(ToolCompletionResponse {
+            content: Some("fixed response".to_string()),
+            tool_calls: vec![],
+            input_tokens: 10,
+            output_tokens: 5,
+            finish_reason: FinishReason::Stop,
+            response_id: None,
+        })
+    }
+
+    fn effective_model_name(&self, _requested_model: Option<&str>) -> String {
+        self.model.to_string()
+    }
+}
+
 // ---------------------------------------------------------------------------
 // Test helpers
 // ---------------------------------------------------------------------------

-async fn start_test_server() -> (SocketAddr, Arc<GatewayState>) {
+async fn start_test_server() -> (SocketAddr, Arc<GatewayState>, Arc<MockLlmState>) {
+    let mock_state = Arc::new(MockLlmState::default());
+
+    let llm_provider: Arc<dyn LlmProvider> = Arc::new(MockLlmProvider::new(mock_state.clone()));
+    let (bound_addr, state) = start_test_server_with_provider(llm_provider).await;
+
+    (bound_addr, state, mock_state)
+}
+
+async fn start_test_server_with_provider(
+    llm_provider: Arc<dyn LlmProvider>,
+) -> (SocketAddr, Arc<GatewayState>) {
    let state = Arc::new(GatewayState {
        msg_tx: tokio::sync::RwLock::new(None),
        sse: SseManager::new(),
@@ -112,7 +198,7 @@ async fn start_test_server() -> (SocketAddr, Arc<GatewayState>) {
        user_id: "test-user".to_string(),
        shutdown_tx: tokio::sync::RwLock::new(None),
        ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
-        llm_provider: Some(Arc::new(MockLlmProvider)),
+        llm_provider: Some(llm_provider),
        skill_registry: None,
        skill_catalog: None,
        chat_rate_limiter: ironclaw::channels::web::server::RateLimiter::new(30, 60),
@@ -139,7 +225,7 @@ fn client() -> reqwest::Client {

 #[tokio::test]
 async fn test_chat_completions_basic() {
-    let (addr, _state) = start_test_server().await;
+    let (addr, _state, mock_state) = start_test_server().await;
    let url = format!("http://{}/v1/chat/completions", addr);

    let resp = client()
@@ -173,11 +259,14 @@ async fn test_chat_completions_basic() {
    assert_eq!(body["usage"]["prompt_tokens"], 10);
    assert_eq!(body["usage"]["completion_tokens"], 5);
    assert_eq!(body["usage"]["total_tokens"], 15);
+
+    let models = mock_state.completion_models.lock().await;
+    assert_eq!(*models, vec![Some("mock-model-v1".to_string())]);
 }

 #[tokio::test]
 async fn test_chat_completions_with_system_message() {
-    let (addr, _state) = start_test_server().await;
+    let (addr, _state, _mock_state) = start_test_server().await;
    let url = format!("http://{}/v1/chat/completions", addr);

    let resp = client()
@@ -204,7 +293,7 @@ async fn test_chat_completions_with_system_message() {

 #[tokio::test]
 async fn test_chat_completions_with_tools() {
-    let (addr, _state) = start_test_server().await;
+    let (addr, _state, mock_state) = start_test_server().await;
    let url = format!("http://{}/v1/chat/completions", addr);

    let resp = client()
@@ -243,11 +332,14 @@ async fn test_chat_completions_with_tools() {
    assert_eq!(tool_calls[0]["id"], "call_mock_001");
    assert_eq!(tool_calls[0]["type"], "function");
    assert_eq!(tool_calls[0]["function"]["name"], "get_weather");
+
+    let models = mock_state.tool_completion_models.lock().await;
+    assert_eq!(*models, vec![Some("mock-model-v1".to_string())]);
 }

 #[tokio::test]
 async fn test_chat_completions_streaming() {
-    let (addr, _state) = start_test_server().await;
+    let (addr, _state, mock_state) = start_test_server().await;
    let url = format!("http://{}/v1/chat/completions", addr);

    let resp = client()
@@ -316,11 +408,14 @@ async fn test_chat_completions_streaming() {
        "Expected reassembled content to contain 'Stream test', got: '{}'",
        full_content
    );
+
+    let models = mock_state.completion_models.lock().await;
+    assert_eq!(*models, vec![Some("mock-model-v1".to_string())]);
 }

 #[tokio::test]
 async fn test_chat_completions_empty_messages() {
-    let (addr, _state) = start_test_server().await;
+    let (addr, _state, _mock_state) = start_test_server().await;
    let url = format!("http://{}/v1/chat/completions", addr);

    let resp = client()
@@ -340,8 +435,8 @@ async fn test_chat_completions_empty_messages() {
 }

 #[tokio::test]
-async fn test_chat_completions_model_mismatch() {
-    let (addr, _state) = start_test_server().await;
+async fn test_chat_completions_model_override() {
+    let (addr, _state, mock_state) = start_test_server().await;
    let url = format!("http://{}/v1/chat/completions", addr);

    let resp = client()
@@ -355,20 +450,173 @@ async fn test_chat_completions_model_mismatch() {
        .await
        .unwrap();

-    assert_eq!(resp.status(), 404);
+    assert_eq!(resp.status(), 200);
+    let body: serde_json::Value = resp.json().await.unwrap();
+    assert_eq!(body["model"], "gpt-4");
+
+    let models = mock_state.completion_models.lock().await;
+    assert_eq!(*models, vec![Some("gpt-4".to_string())]);
+}
+
+#[tokio::test]
+async fn test_chat_completions_uses_effective_model_when_override_ignored() {
+    let provider: Arc<dyn LlmProvider> = Arc::new(FixedModelProvider::new("configured-model"));
+    let (addr, _state) = start_test_server_with_provider(provider).await;
+    let url = format!("http://{}/v1/chat/completions", addr);
+
+    let resp = client()
+        .post(&url)
+        .bearer_auth(AUTH_TOKEN)
+        .json(&serde_json::json!({
+            "model": "gpt-4",
+            "messages": [{"role": "user", "content": "Hi"}]
+        }))
+        .send()
+        .await
+        .unwrap();
+
+    assert_eq!(resp.status(), 200);
+    let body: serde_json::Value = resp.json().await.unwrap();
+    assert_eq!(body["model"], "configured-model");
+}
+
+#[tokio::test]
+async fn test_chat_completions_streaming_uses_effective_model_when_override_ignored() {
+    let provider: Arc<dyn LlmProvider> = Arc::new(FixedModelProvider::new("configured-model"));
+    let (addr, _state) = start_test_server_with_provider(provider).await;
+    let url = format!("http://{}/v1/chat/completions", addr);
+
+    let resp = client()
+        .post(&url)
+        .bearer_auth(AUTH_TOKEN)
+        .json(&serde_json::json!({
+            "model": "gpt-4",
+            "messages": [{"role": "user", "content": "Hi"}],
+            "stream": true
+        }))
+        .send()
+        .await
+        .unwrap();
+
+    assert_eq!(resp.status(), 200);
+    let text = resp.text().await.unwrap();
+    assert!(
+        text.contains("\"model\":\"configured-model\""),
+        "Expected streaming chunks to report configured model, got: {}",
+        text
+    );
+}
+
+#[tokio::test]
+async fn test_chat_completions_model_too_long() {
+    let (addr, _state, mock_state) = start_test_server().await;
+    let url = format!("http://{}/v1/chat/completions", addr);
+
+    let resp = client()
+        .post(&url)
+        .bearer_auth(AUTH_TOKEN)
+        .json(&serde_json::json!({
+            "model": "m".repeat(300),
+            "messages": [{"role": "user", "content": "Hi"}]
+        }))
+        .send()
+        .await
+        .unwrap();
+
+    assert_eq!(resp.status(), 400);
    let body: serde_json::Value = resp.json().await.unwrap();
-    assert_eq!(body["error"]["code"], "model_not_found");
    assert!(
        body["error"]["message"]
            .as_str()
-            .unwrap()
-            .contains("mock-model-v1")
+            .unwrap_or("")
+            .contains("model"),
+        "Expected model validation error, got: {}",
+        body
+    );
+
+    // Validation should fail before provider invocation.
+    let models = mock_state.completion_models.lock().await;
+    assert!(
+        models.is_empty(),
+        "provider should not be called: {:?}",
+        *models
+    );
+}
+
+#[tokio::test]
+async fn test_chat_completions_model_with_control_chars() {
+    let (addr, _state, mock_state) = start_test_server().await;
+    let url = format!("http://{}/v1/chat/completions", addr);
+
+    let resp = client()
+        .post(&url)
+        .bearer_auth(AUTH_TOKEN)
+        .json(&serde_json::json!({
+            "model": "gpt-4\noops",
+            "messages": [{"role": "user", "content": "Hi"}]
+        }))
+        .send()
+        .await
+        .unwrap();
+
+    assert_eq!(resp.status(), 400);
+    let body: serde_json::Value = resp.json().await.unwrap();
+    assert!(
+        body["error"]["message"]
+            .as_str()
+            .unwrap_or("")
+            .contains("control"),
+        "Expected model validation error, got: {}",
+        body
+    );
+
+    // Validation should fail before provider invocation.
+    let models = mock_state.completion_models.lock().await;
+    assert!(
+        models.is_empty(),
+        "provider should not be called: {:?}",
+        *models
+    );
+}
+
+#[tokio::test]
+async fn test_chat_completions_model_with_surrounding_whitespace() {
+    let (addr, _state, mock_state) = start_test_server().await;
+    let url = format!("http://{}/v1/chat/completions", addr);
+
+    let resp = client()
+        .post(&url)
+        .bearer_auth(AUTH_TOKEN)
+        .json(&serde_json::json!({
+            "model": " gpt-4 ",
+            "messages": [{"role": "user", "content": "Hi"}]
+        }))
+        .send()
+        .await
+        .unwrap();
+
+    assert_eq!(resp.status(), 400);
+    let body: serde_json::Value = resp.json().await.unwrap();
+    assert!(
+        body["error"]["message"]
+            .as_str()
+            .unwrap_or("")
+            .contains("leading or trailing whitespace"),
+        "Expected model validation error, got: {}",
+        body
+    );
+
+    let models = mock_state.completion_models.lock().await;
+    assert!(
+        models.is_empty(),
+        "provider should not be called: {:?}",
+        *models
    );
 }

 #[tokio::test]
 async fn test_chat_completions_no_auth() {
-    let (addr, _state) = start_test_server().await;
+    let (addr, _state, _mock_state) = start_test_server().await;
    let url = format!("http://{}/v1/chat/completions", addr);

    let resp = client()
@@ -387,7 +635,7 @@ async fn test_chat_completions_no_auth() {

 #[tokio::test]
 async fn test_models_endpoint() {
-    let (addr, _state) = start_test_server().await;
+    let (addr, _state, _mock_state) = start_test_server().await;
    let url = format!("http://{}/v1/models", addr);

    let resp = client()
@@ -410,7 +658,7 @@ async fn test_models_endpoint() {

 #[tokio::test]
 async fn test_models_no_auth() {
-    let (addr, _state) = start_test_server().await;
+    let (addr, _state, _mock_state) = start_test_server().await;
    let url = format!("http://{}/v1/models", addr);

    let resp = client().get(&url).send().await.unwrap();
@@ -462,7 +710,7 @@ async fn test_no_llm_provider_returns_503() {

 #[tokio::test]
 async fn test_chat_completions_body_too_large() {
-    let (addr, _state) = start_test_server().await;
+    let (addr, _state, _mock_state) = start_test_server().await;
    let url = format!("http://{}/v1/chat/completions", addr);

    // Build a payload over 1 MB (the gateway's DefaultBodyLimit)