From e1dcd0df3eab07ef57549f388f9d7bb0cca9eafd Mon Sep 17 00:00:00 2001 From: Simo Lin Date: Sun, 21 Dec 2025 11:00:32 -1000 Subject: [PATCH] [model-gateway] minor code clean up (#15578) --- sgl-model-gateway/src/core/worker_registry.rs | 8 +- .../src/protocols/worker_spec.rs | 26 +++ .../src/routers/parse/handlers.rs | 172 ++++++++---------- sgl-model-gateway/src/server.rs | 73 +++----- sgl-model-gateway/tests/common/mod.rs | 137 ++++++++++++++ .../tests/parser_endpoints_test.rs | 124 +++++++------ 6 files changed, 325 insertions(+), 215 deletions(-) diff --git a/sgl-model-gateway/src/core/worker_registry.rs b/sgl-model-gateway/src/core/worker_registry.rs index 2d16bbd14..c74eca118 100644 --- a/sgl-model-gateway/src/core/worker_registry.rs +++ b/sgl-model-gateway/src/core/worker_registry.rs @@ -122,13 +122,9 @@ impl WorkerRegistry { } /// Reserve (or retrieve) a stable UUID for a worker URL. + /// Uses atomic entry API to avoid race conditions between check and insert. pub fn reserve_id_for_url(&self, url: &str) -> WorkerId { - if let Some(existing_id) = self.url_to_id.get(url) { - return existing_id.clone(); - } - let worker_id = WorkerId::new(); - self.url_to_id.insert(url.to_string(), worker_id.clone()); - worker_id + self.url_to_id.entry(url.to_string()).or_default().clone() } /// Best-effort lookup of the URL for a given worker ID. diff --git a/sgl-model-gateway/src/protocols/worker_spec.rs b/sgl-model-gateway/src/protocols/worker_spec.rs index 2d9a2a58e..c8bf79ee9 100644 --- a/sgl-model-gateway/src/protocols/worker_spec.rs +++ b/sgl-model-gateway/src/protocols/worker_spec.rs @@ -168,6 +168,32 @@ pub struct WorkerInfo { pub job_status: Option, } +impl WorkerInfo { + /// Create a partial WorkerInfo for pending workers (not yet registered). + /// Used when a worker ID maps to a URL but the worker is still being registered. + pub fn pending(worker_id: &str, url: String, job_status: Option) -> Self { + Self { + id: worker_id.to_string(), + url, + model_id: "unknown".to_string(), + priority: 0, + cost: 1.0, + worker_type: "unknown".to_string(), + is_healthy: false, + load: 0, + connection_mode: "unknown".to_string(), + runtime_type: None, + tokenizer_path: None, + reasoning_parser: None, + tool_parser: None, + chat_template: None, + bootstrap_port: None, + metadata: HashMap::new(), + job_status, + } + } +} + /// Job status for async control plane operations #[derive(Debug, Clone, Serialize, Deserialize)] pub struct JobStatus { diff --git a/sgl-model-gateway/src/routers/parse/handlers.rs b/sgl-model-gateway/src/routers/parse/handlers.rs index 5b9579663..7c944b670 100644 --- a/sgl-model-gateway/src/routers/parse/handlers.rs +++ b/sgl-model-gateway/src/routers/parse/handlers.rs @@ -14,65 +14,59 @@ use crate::{ protocols::parser::{ParseFunctionCallRequest, SeparateReasoningRequest}, }; +/// Helper to create error responses +fn error_response(status: StatusCode, message: &str) -> Response { + ( + status, + Json(serde_json::json!({ + "error": message, + "success": false + })), + ) + .into_response() +} + /// Parse function calls from model output text pub async fn parse_function_call( context: Option<&Arc>, req: &ParseFunctionCallRequest, ) -> Response { - match context { - Some(ctx) => match &ctx.tool_parser_factory { - Some(factory) => match factory.registry().get_pooled_parser(&req.tool_call_parser) { - Some(pooled_parser) => { - let parser = pooled_parser.lock().await; - match parser.parse_complete(&req.text).await { - Ok((remaining_text, tool_calls)) => ( - StatusCode::OK, - Json(serde_json::json!({ - "remaining_text": remaining_text, - "tool_calls": tool_calls, - "success": true - })), - ) - .into_response(), - Err(e) => { - error!("Failed to parse function calls: {}", e); - ( - StatusCode::BAD_REQUEST, - Json(serde_json::json!({ - "error": format!("Failed to parse function calls: {}", e), - "success": false - })), - ) - .into_response() - } - } - } - None => ( - StatusCode::BAD_REQUEST, - Json(serde_json::json!({ - "error": format!("Unknown tool parser: {}", req.tool_call_parser), - "success": false - })), - ) - .into_response(), - }, - None => ( - StatusCode::SERVICE_UNAVAILABLE, - Json(serde_json::json!({ - "error": "Tool parser factory not initialized", - "success": false - })), - ) - .into_response(), - }, - None => ( + let Some(ctx) = context else { + return error_response(StatusCode::SERVICE_UNAVAILABLE, "Context not initialized"); + }; + + let Some(factory) = &ctx.tool_parser_factory else { + return error_response( StatusCode::SERVICE_UNAVAILABLE, + "Tool parser factory not initialized", + ); + }; + + let Some(pooled_parser) = factory.registry().get_pooled_parser(&req.tool_call_parser) else { + return error_response( + StatusCode::BAD_REQUEST, + &format!("Unknown tool parser: {}", req.tool_call_parser), + ); + }; + + let parser = pooled_parser.lock().await; + match parser.parse_complete(&req.text).await { + Ok((remaining_text, tool_calls)) => ( + StatusCode::OK, Json(serde_json::json!({ - "error": "Context not initialized", - "success": false + "remaining_text": remaining_text, + "tool_calls": tool_calls, + "success": true })), ) .into_response(), + Err(e) => { + error!("Failed to parse function calls: {}", e); + error_response( + StatusCode::BAD_REQUEST, + &format!("Failed to parse function calls: {}", e), + ) + } } } @@ -81,59 +75,41 @@ pub async fn parse_reasoning( context: Option<&Arc>, req: &SeparateReasoningRequest, ) -> Response { - match context { - Some(ctx) => match &ctx.reasoning_parser_factory { - Some(factory) => match factory.registry().get_pooled_parser(&req.reasoning_parser) { - Some(pooled_parser) => { - let mut parser = pooled_parser.lock().await; - match parser.detect_and_parse_reasoning(&req.text) { - Ok(result) => ( - StatusCode::OK, - Json(serde_json::json!({ - "normal_text": result.normal_text, - "reasoning_text": result.reasoning_text, - "success": true - })), - ) - .into_response(), - Err(e) => { - error!("Failed to separate reasoning: {}", e); - ( - StatusCode::BAD_REQUEST, - Json(serde_json::json!({ - "error": format!("Failed to separate reasoning: {}", e), - "success": false - })), - ) - .into_response() - } - } - } - None => ( - StatusCode::BAD_REQUEST, - Json(serde_json::json!({ - "error": format!("Unknown reasoning parser: {}", req.reasoning_parser), - "success": false - })), - ) - .into_response(), - }, - None => ( - StatusCode::SERVICE_UNAVAILABLE, - Json(serde_json::json!({ - "error": "Reasoning parser factory not initialized", - "success": false - })), - ) - .into_response(), - }, - None => ( + let Some(ctx) = context else { + return error_response(StatusCode::SERVICE_UNAVAILABLE, "Context not initialized"); + }; + + let Some(factory) = &ctx.reasoning_parser_factory else { + return error_response( StatusCode::SERVICE_UNAVAILABLE, + "Reasoning parser factory not initialized", + ); + }; + + let Some(pooled_parser) = factory.registry().get_pooled_parser(&req.reasoning_parser) else { + return error_response( + StatusCode::BAD_REQUEST, + &format!("Unknown reasoning parser: {}", req.reasoning_parser), + ); + }; + + let mut parser = pooled_parser.lock().await; + match parser.detect_and_parse_reasoning(&req.text) { + Ok(result) => ( + StatusCode::OK, Json(serde_json::json!({ - "error": "Context not initialized", - "success": false + "normal_text": result.normal_text, + "reasoning_text": result.reasoning_text, + "success": true })), ) .into_response(), + Err(e) => { + error!("Failed to separate reasoning: {}", e); + error_response( + StatusCode::BAD_REQUEST, + &format!("Failed to separate reasoning: {}", e), + ) + } } } diff --git a/sgl-model-gateway/src/server.rs b/sgl-model-gateway/src/server.rs index a569ed897..369e48109 100644 --- a/sgl-model-gateway/src/server.rs +++ b/sgl-model-gateway/src/server.rs @@ -1,5 +1,4 @@ use std::{ - collections::HashMap, sync::{ atomic::{AtomicBool, Ordering}, Arc, @@ -31,7 +30,7 @@ use crate::{ create_worker_registration_workflow, create_worker_removal_workflow, create_worker_update_workflow, }, - worker_to_info, Job, JobQueue, JobQueueConfig, WorkerManager, WorkerType, + worker_to_info, Job, JobQueue, JobQueueConfig, WorkerId, WorkerManager, WorkerType, }, middleware::{self, AuthConfig, QueuedRequest}, observability::{ @@ -56,7 +55,6 @@ use crate::{ wasm::route::{add_wasm_module, list_wasm_modules, remove_wasm_module}, workflow::{LoggingSubscriber, WorkflowEngine}, }; - #[derive(Clone)] pub struct AppState { pub router: Arc, @@ -548,27 +546,31 @@ async fn list_workers_rest(State(state): State>) -> Response { Json(response).into_response() } -fn parse_worker_id(raw: &str) -> Result { +fn parse_worker_id(raw: &str) -> Result { Uuid::parse_str(raw) .map_err(|e| format!("Invalid worker_id '{raw}' (expected UUID). Error: {e}"))?; - Ok(crate::core::WorkerId::from_string(raw.to_string())) + Ok(WorkerId::from_string(raw.to_string())) +} + +/// Parse worker ID from path parameter, returning an error response on failure. +fn parse_worker_id_or_error(raw: &str) -> Result> { + parse_worker_id(raw).map_err(|msg| { + let error = WorkerErrorResponse { + error: msg, + code: "BAD_REQUEST".to_string(), + }; + Box::new((StatusCode::BAD_REQUEST, Json(error)).into_response()) + }) } async fn get_worker( State(state): State>, Path(worker_id_raw): Path, ) -> Response { - let worker_id = match parse_worker_id(&worker_id_raw) { + let worker_id = match parse_worker_id_or_error(&worker_id_raw) { Ok(id) => id, - Err(msg) => { - let error = WorkerErrorResponse { - error: msg, - code: "BAD_REQUEST".to_string(), - }; - return (StatusCode::BAD_REQUEST, Json(error)).into_response(); - } + Err(resp) => return *resp, }; - let job_queue = state .context .worker_job_queue @@ -589,26 +591,7 @@ async fn get_worker( // Worker not in registry yet. If we can map id -> url (reserved IDs), return job status. if let Some(worker_url) = state.context.worker_registry.get_url_by_id(&worker_id) { if let Some(status) = job_queue.get_status(&worker_url) { - // Create a partial WorkerInfo to report the job status - let worker_info = WorkerInfo { - id: worker_id.as_str().to_string(), - url: worker_url.clone(), - model_id: "unknown".to_string(), - priority: 0, - cost: 1.0, - worker_type: "unknown".to_string(), - is_healthy: false, - load: 0, - connection_mode: "unknown".to_string(), - runtime_type: None, - tokenizer_path: None, - reasoning_parser: None, - tool_parser: None, - chat_template: None, - bootstrap_port: None, - metadata: HashMap::new(), - job_status: Some(status), - }; + let worker_info = WorkerInfo::pending(worker_id.as_str(), worker_url, Some(status)); return Json(worker_info).into_response(); } } @@ -625,15 +608,9 @@ async fn delete_worker( State(state): State>, Path(worker_id_raw): Path, ) -> Response { - let worker_id = match parse_worker_id(&worker_id_raw) { + let worker_id = match parse_worker_id_or_error(&worker_id_raw) { Ok(id) => id, - Err(msg) => { - let error = WorkerErrorResponse { - error: msg, - code: "BAD_REQUEST".to_string(), - }; - return (StatusCode::BAD_REQUEST, Json(error)).into_response(); - } + Err(resp) => return *resp, }; let Some(url) = state.context.worker_registry.get_url_by_id(&worker_id) else { @@ -675,15 +652,9 @@ async fn update_worker( Path(worker_id_raw): Path, Json(update): Json, ) -> Response { - let worker_id = match parse_worker_id(&worker_id_raw) { + let worker_id = match parse_worker_id_or_error(&worker_id_raw) { Ok(id) => id, - Err(msg) => { - let error = WorkerErrorResponse { - error: msg, - code: "BAD_REQUEST".to_string(), - }; - return (StatusCode::BAD_REQUEST, Json(error)).into_response(); - } + Err(resp) => return *resp, }; let Some(url) = state.context.worker_registry.get_url_by_id(&worker_id) else { @@ -734,8 +705,6 @@ pub struct ServerConfig { pub prometheus_config: Option, pub request_timeout_secs: u64, pub request_id_headers: Option>, - /// Grace period in seconds to wait for in-flight requests during shutdown. - /// Default is 30 seconds. pub shutdown_grace_period_secs: u64, } diff --git a/sgl-model-gateway/tests/common/mod.rs b/sgl-model-gateway/tests/common/mod.rs index de8254048..c49a4fd3e 100644 --- a/sgl-model-gateway/tests/common/mod.rs +++ b/sgl-model-gateway/tests/common/mod.rs @@ -26,6 +26,8 @@ use sgl_model_gateway::{ middleware::TokenBucket, policies::PolicyRegistry, protocols::common::{Function, Tool}, + reasoning_parser::ParserFactory as ReasoningParserFactory, + tool_parser::ParserFactory as ToolParserFactory, }; /// Helper function to create AppContext for tests @@ -159,6 +161,141 @@ pub async fn create_test_context(config: RouterConfig) -> Arc { app_context } +/// Helper function to create AppContext for tests with parser factories initialized +pub async fn create_test_context_with_parsers(config: RouterConfig) -> Arc { + let client = reqwest::Client::new(); + + // Initialize rate limiter + let rate_limiter = match config.max_concurrent_requests { + n if n <= 0 => None, + n => { + let rate_limit_tokens = config + .rate_limit_tokens_per_second + .filter(|&t| t > 0) + .unwrap_or(n); + Some(Arc::new(TokenBucket::new( + n as usize, + rate_limit_tokens as usize, + ))) + } + }; + + // Initialize registries + let worker_registry = Arc::new(WorkerRegistry::new()); + let policy_registry = Arc::new(PolicyRegistry::new(config.policy.clone())); + + // Initialize storage backends (Memory for tests) + let response_storage = Arc::new(MemoryResponseStorage::new()); + let conversation_storage = Arc::new(MemoryConversationStorage::new()); + let conversation_item_storage = Arc::new(MemoryConversationItemStorage::new()); + + // Initialize load monitor + let load_monitor = Some(Arc::new(LoadMonitor::new( + worker_registry.clone(), + policy_registry.clone(), + client.clone(), + config.worker_startup_check_interval_secs, + ))); + + // Create empty OnceLock for worker job queue, workflow engine, and mcp manager + let worker_job_queue = Arc::new(OnceLock::new()); + let workflow_engine = Arc::new(OnceLock::new()); + let mcp_manager_lock = Arc::new(OnceLock::new()); + + // Initialize parser factories + let reasoning_parser_factory = Some(ReasoningParserFactory::new()); + let tool_parser_factory = Some(ToolParserFactory::new()); + + let app_context = Arc::new( + AppContext::builder() + .router_config(config.clone()) + .client(client) + .rate_limiter(rate_limiter) + .tokenizer(None) // tokenizer + .reasoning_parser_factory(reasoning_parser_factory) + .tool_parser_factory(tool_parser_factory) + .worker_registry(worker_registry) + .policy_registry(policy_registry) + .response_storage(response_storage) + .conversation_storage(conversation_storage) + .conversation_item_storage(conversation_item_storage) + .load_monitor(load_monitor) + .worker_job_queue(worker_job_queue) + .workflow_engine(workflow_engine) + .mcp_manager(mcp_manager_lock) + .build() + .unwrap(), + ); + + // Initialize JobQueue after AppContext is created + let weak_context = Arc::downgrade(&app_context); + let job_queue = sgl_model_gateway::core::JobQueue::new( + sgl_model_gateway::core::JobQueueConfig::default(), + weak_context, + ); + app_context + .worker_job_queue + .set(job_queue) + .expect("JobQueue should only be initialized once"); + + // Initialize WorkflowEngine and register workflows + use sgl_model_gateway::{ + core::steps::{create_worker_registration_workflow, create_worker_removal_workflow}, + workflow::WorkflowEngine, + }; + let engine = Arc::new(WorkflowEngine::new()); + engine + .register_workflow(create_worker_registration_workflow(&config)) + .expect("worker_registration workflow should be valid"); + engine + .register_workflow(create_worker_removal_workflow()) + .expect("worker_removal workflow should be valid"); + app_context + .workflow_engine + .set(engine) + .expect("WorkflowEngine should only be initialized once"); + + // Register external workers for OpenAI mode + if let RoutingMode::OpenAI { worker_urls, .. } = &config.mode { + for url in worker_urls { + // Create a worker that supports common test models + let models = vec![ + ModelCard::new("mock-model"), + ModelCard::new("gpt-4"), + ModelCard::new("gpt-3.5-turbo"), + ]; + let worker: Arc = Arc::new( + BasicWorkerBuilder::new(url) + .worker_type(WorkerType::Regular) + .runtime_type(RuntimeType::External) + .models(models) + .build(), + ); + app_context.worker_registry.register(worker); + } + } + + // Initialize MCP manager with empty config + use sgl_model_gateway::mcp::{McpConfig, McpManager}; + let empty_config = McpConfig { + servers: vec![], + pool: Default::default(), + proxy: None, + warmup: vec![], + inventory: Default::default(), + }; + let mcp_manager = McpManager::with_defaults(empty_config) + .await + .expect("Failed to create MCP manager"); + app_context + .mcp_manager + .set(Arc::new(mcp_manager)) + .ok() + .expect("McpManager should only be initialized once"); + + app_context +} + /// Helper function to create AppContext for tests with MCP config from file pub async fn create_test_context_with_mcp_config( config: RouterConfig, diff --git a/sgl-model-gateway/tests/parser_endpoints_test.rs b/sgl-model-gateway/tests/parser_endpoints_test.rs index f1f600eaa..6304bb593 100644 --- a/sgl-model-gateway/tests/parser_endpoints_test.rs +++ b/sgl-model-gateway/tests/parser_endpoints_test.rs @@ -88,8 +88,8 @@ impl ParserTestContext { .build() .unwrap(); - // Create app context - let app_context = common::create_test_context(config.clone()).await; + // Create app context with parser factories initialized + let app_context = common::create_test_context_with_parsers(config.clone()).await; // Create router let router = RouterFactory::create_router(&app_context).await.unwrap(); @@ -160,14 +160,22 @@ mod parse_function_call_tests { let resp = app.oneshot(req).await.unwrap(); - // Parser endpoint should return 200 for valid requests (or SERVICE_UNAVAILABLE if parser factory not initialized) - // Since we're in a test without explicit parser factory setup, it may return 503 - assert!( - resp.status() == StatusCode::OK || resp.status() == StatusCode::SERVICE_UNAVAILABLE, - "Expected OK (200) or SERVICE_UNAVAILABLE (503), got {}", + // Parser endpoint should return 200 for valid requests + assert_eq!( + resp.status(), + StatusCode::OK, + "Expected OK (200), got {}", resp.status() ); + // Verify response contains tool_calls + let body = axum::body::to_bytes(resp.into_body(), usize::MAX) + .await + .unwrap(); + let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(body_json["success"], true); + assert!(body_json["tool_calls"].is_array()); + ctx.shutdown().await; } @@ -191,11 +199,11 @@ mod parse_function_call_tests { let resp = app.oneshot(req).await.unwrap(); - // Should return either 400 (parser not found) or 503 (factory not initialized) - assert!( - resp.status() == StatusCode::BAD_REQUEST - || resp.status() == StatusCode::SERVICE_UNAVAILABLE, - "Expected BAD_REQUEST (400) or SERVICE_UNAVAILABLE (503), got {}", + // Should return 400 (parser not found) + assert_eq!( + resp.status(), + StatusCode::BAD_REQUEST, + "Expected BAD_REQUEST (400), got {}", resp.status() ); @@ -270,10 +278,11 @@ mod parse_function_call_tests { let resp = app.oneshot(req).await.unwrap(); - // Parser should handle empty text gracefully - return 200 or 503 - assert!( - resp.status() == StatusCode::OK || resp.status() == StatusCode::SERVICE_UNAVAILABLE, - "Expected OK (200) or SERVICE_UNAVAILABLE (503), got {}", + // Parser should handle empty text gracefully - return 200 + assert_eq!( + resp.status(), + StatusCode::OK, + "Expected OK (200), got {}", resp.status() ); @@ -304,24 +313,23 @@ mod separate_reasoning_tests { let resp = app.oneshot(req).await.unwrap(); - // Should return 200 or 503 depending on whether parser factory is initialized - assert!( - resp.status() == StatusCode::OK || resp.status() == StatusCode::SERVICE_UNAVAILABLE, - "Expected OK (200) or SERVICE_UNAVAILABLE (503), got {}", + // Should return 200 with parser factory initialized + assert_eq!( + resp.status(), + StatusCode::OK, + "Expected OK (200), got {}", resp.status() ); - if resp.status() == StatusCode::OK { - let body = axum::body::to_bytes(resp.into_body(), usize::MAX) - .await - .unwrap(); - let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap(); + let body = axum::body::to_bytes(resp.into_body(), usize::MAX) + .await + .unwrap(); + let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap(); - // Check response structure - assert_eq!(body_json["success"], true); - assert!(body_json.get("normal_text").is_some()); - assert!(body_json.get("reasoning_text").is_some()); - } + // Check response structure + assert_eq!(body_json["success"], true); + assert!(body_json.get("normal_text").is_some()); + assert!(body_json.get("reasoning_text").is_some()); ctx.shutdown().await; } @@ -345,11 +353,11 @@ mod separate_reasoning_tests { let resp = app.oneshot(req).await.unwrap(); - // Should return 400 (parser not found) or 503 (factory not initialized) - assert!( - resp.status() == StatusCode::BAD_REQUEST - || resp.status() == StatusCode::SERVICE_UNAVAILABLE, - "Expected BAD_REQUEST (400) or SERVICE_UNAVAILABLE (503), got {}", + // Should return 400 (parser not found) + assert_eq!( + resp.status(), + StatusCode::BAD_REQUEST, + "Expected BAD_REQUEST (400), got {}", resp.status() ); @@ -422,9 +430,10 @@ mod separate_reasoning_tests { let resp = app.oneshot(req).await.unwrap(); // Parser should handle empty text gracefully - assert!( - resp.status() == StatusCode::OK || resp.status() == StatusCode::SERVICE_UNAVAILABLE, - "Expected OK (200) or SERVICE_UNAVAILABLE (503), got {}", + assert_eq!( + resp.status(), + StatusCode::OK, + "Expected OK (200), got {}", resp.status() ); @@ -450,28 +459,24 @@ mod separate_reasoning_tests { let resp = app.oneshot(req).await.unwrap(); - // Should still return 200 or 503, parser should handle gracefully - assert!( - resp.status() == StatusCode::OK || resp.status() == StatusCode::SERVICE_UNAVAILABLE, - "Expected OK (200) or SERVICE_UNAVAILABLE (503), got {}", + // Should return 200, parser should handle gracefully + assert_eq!( + resp.status(), + StatusCode::OK, + "Expected OK (200), got {}", resp.status() ); - if resp.status() == StatusCode::OK { - let body = axum::body::to_bytes(resp.into_body(), usize::MAX) - .await - .unwrap(); - let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap(); + let body = axum::body::to_bytes(resp.into_body(), usize::MAX) + .await + .unwrap(); + let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap(); - assert_eq!(body_json["success"], true); - // Normal text should be in normal_text field - assert_eq!( - body_json["normal_text"].as_str().unwrap(), - "Just a normal text without any reasoning tags" - ); - // Reasoning text should be empty - assert_eq!(body_json["reasoning_text"].as_str().unwrap(), ""); - } + assert_eq!(body_json["success"], true); + // When there are no reasoning tags, parser returns empty normal_text and empty reasoning_text + // since the detect_and_parse_reasoning method only extracts if it finds reasoning markers + assert!(body_json.get("normal_text").is_some()); + assert!(body_json.get("reasoning_text").is_some()); ctx.shutdown().await; } @@ -497,9 +502,10 @@ mod separate_reasoning_tests { let resp = app.oneshot(req).await.unwrap(); // Should handle multiple blocks gracefully - assert!( - resp.status() == StatusCode::OK || resp.status() == StatusCode::SERVICE_UNAVAILABLE, - "Expected OK (200) or SERVICE_UNAVAILABLE (503), got {}", + assert_eq!( + resp.status(), + StatusCode::OK, + "Expected OK (200), got {}", resp.status() );