chore: vendor sglang v0.5.10 snapshot

third_party/sglang/sgl-model-gateway/tests/common/test_app.rs

@@ -0,0 +1,267 @@
+use std::sync::{Arc, OnceLock};
+
+use axum::Router;
+use data_connector::{
+    MemoryConversationItemStorage, MemoryConversationStorage, MemoryResponseStorage,
+};
+use reqwest::Client;
+use smg::{
+    app_context::AppContext,
+    config::RouterConfig,
+    core::{
+        BasicWorkerBuilder, LoadMonitor, ModelCard, RuntimeType, Worker, WorkerRegistry, WorkerType,
+    },
+    middleware::{AuthConfig, TokenBucket},
+    policies::PolicyRegistry,
+    routers::RouterTrait,
+    server::{build_app, AppState},
+    tokenizer::registry::TokenizerRegistry,
+};
+use smg_mcp::{McpConfig, McpManager};
+
+/// Create a test Axum application using the actual server's build_app function
+#[allow(dead_code)]
+pub fn create_test_app(
+    router: Arc<dyn RouterTrait>,
+    client: Client,
+    router_config: &RouterConfig,
+) -> Router {
+    // Initialize rate limiter
+    let rate_limiter = match router_config.max_concurrent_requests {
+        n if n <= 0 => None,
+        n => {
+            let rate_limit_tokens = router_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(router_config.policy.clone()));
+
+    // Initialize storage backends
+    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(),
+        router_config.worker_startup_check_interval_secs,
+    )));
+
+    // Create empty OnceLock for worker job queue and workflow engines
+    let worker_job_queue = Arc::new(OnceLock::new());
+    let workflow_engines = Arc::new(OnceLock::new());
+
+    // Create AppContext using builder pattern
+    let app_context = Arc::new(
+        AppContext::builder()
+            .router_config(router_config.clone())
+            .client(client)
+            .rate_limiter(rate_limiter)
+            .tokenizer_registry(Arc::new(TokenizerRegistry::new())) // tokenizer
+            .reasoning_parser_factory(None) // reasoning_parser_factory
+            .tool_parser_factory(None) // 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_engines(workflow_engines)
+            .build()
+            .unwrap(),
+    );
+
+    // Create AppState with the test router and context
+    let app_state = Arc::new(AppState {
+        router,
+        context: app_context,
+        concurrency_queue_tx: None,
+        router_manager: None,
+        mesh_handler: None,
+        mesh_sync_manager: None,
+    });
+
+    // Configure request ID headers (use defaults if not specified)
+    let request_id_headers = router_config.request_id_headers.clone().unwrap_or_else(|| {
+        vec![
+            "x-request-id".to_string(),
+            "x-correlation-id".to_string(),
+            "x-trace-id".to_string(),
+            "request-id".to_string(),
+        ]
+    });
+
+    // Create auth config from router config
+    let auth_config = AuthConfig {
+        api_key: router_config.api_key.clone(),
+    };
+
+    // Use the actual server's build_app function
+    build_app(
+        app_state,
+        auth_config,
+        None, // No control plane auth for tests
+        router_config.max_payload_size,
+        request_id_headers,
+        router_config.cors_allowed_origins.clone(),
+    )
+}
+
+/// Create a test Axum application with an existing AppContext
+#[allow(dead_code)]
+pub fn create_test_app_with_context(
+    router: Arc<dyn RouterTrait>,
+    app_context: Arc<AppContext>,
+) -> Router {
+    // Create AppState with the test router and context
+    let app_state = Arc::new(AppState {
+        router,
+        context: app_context.clone(),
+        concurrency_queue_tx: None,
+        router_manager: None,
+        mesh_handler: None,
+        mesh_sync_manager: None,
+    });
+
+    // Get config from the context
+    let router_config = &app_context.router_config;
+
+    // Configure request ID headers (use defaults if not specified)
+    let request_id_headers = router_config.request_id_headers.clone().unwrap_or_else(|| {
+        vec![
+            "x-request-id".to_string(),
+            "x-correlation-id".to_string(),
+            "x-trace-id".to_string(),
+            "request-id".to_string(),
+        ]
+    });
+
+    // Create auth config from router config
+    let auth_config = AuthConfig {
+        api_key: router_config.api_key.clone(),
+    };
+
+    // Use the actual server's build_app function
+    build_app(
+        app_state,
+        auth_config,
+        None, // No control plane auth for tests
+        router_config.max_payload_size,
+        request_id_headers,
+        router_config.cors_allowed_origins.clone(),
+    )
+}
+
+/// Create a minimal test AppContext for unit tests
+#[allow(dead_code)]
+pub async fn create_test_app_context() -> Arc<AppContext> {
+    let router_config = RouterConfig::default();
+    let client = Client::new();
+
+    // Initialize empty OnceLocks
+    let worker_job_queue = Arc::new(OnceLock::new());
+    let workflow_engines = Arc::new(OnceLock::new());
+
+    // Initialize MCP manager with empty config
+    let mcp_manager_lock = Arc::new(OnceLock::new());
+    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");
+    mcp_manager_lock.set(Arc::new(mcp_manager)).ok();
+
+    // Initialize registries
+    let worker_registry = Arc::new(WorkerRegistry::new());
+    let policy_registry = Arc::new(PolicyRegistry::new(router_config.policy.clone()));
+
+    // Initialize storage backends
+    let response_storage = Arc::new(MemoryResponseStorage::new());
+    let conversation_storage = Arc::new(MemoryConversationStorage::new());
+    let conversation_item_storage = Arc::new(MemoryConversationItemStorage::new());
+
+    Arc::new(
+        AppContext::builder()
+            .router_config(router_config)
+            .client(client)
+            .rate_limiter(None)
+            .tokenizer_registry(Arc::new(TokenizerRegistry::new()))
+            .reasoning_parser_factory(None)
+            .tool_parser_factory(None)
+            .worker_registry(worker_registry)
+            .policy_registry(policy_registry)
+            .response_storage(response_storage)
+            .conversation_storage(conversation_storage)
+            .conversation_item_storage(conversation_item_storage)
+            .load_monitor(None)
+            .worker_job_queue(worker_job_queue)
+            .workflow_engines(workflow_engines)
+            .mcp_manager(mcp_manager_lock)
+            .build()
+            .unwrap(),
+    )
+}
+
+/// Register an external worker (OpenAI-compatible API endpoint) in the test AppContext.
+///
+/// This is used by tests that need to test the OpenAI router, which expects
+/// workers to be registered in the WorkerRegistry before routing requests.
+///
+/// # Arguments
+/// * `ctx` - The AppContext to register the worker in
+/// * `url` - The base URL of the external API endpoint
+/// * `models` - Optional list of model IDs this worker supports. If empty, uses "gpt-3.5-turbo" as default.
+#[allow(dead_code)]
+pub fn register_external_worker(ctx: &Arc<AppContext>, url: &str, models: Option<Vec<&str>>) {
+    let model_list: Vec<ModelCard> = models
+        .unwrap_or_else(|| vec!["gpt-3.5-turbo"])
+        .into_iter()
+        .map(ModelCard::new)
+        .collect();
+
+    let worker: Arc<dyn Worker> = Arc::new(
+        BasicWorkerBuilder::new(url)
+            .worker_type(WorkerType::Regular)
+            .runtime_type(RuntimeType::External)
+            .models(model_list)
+            .build(),
+    );
+
+    ctx.worker_registry.register(worker);
+}
+
+/// Register an external worker with a custom model card that has aliases.
+///
+/// # Arguments
+/// * `ctx` - The AppContext to register the worker in
+/// * `url` - The base URL of the external API endpoint
+/// * `model_card` - A fully configured ModelCard with aliases, provider, etc.
+#[allow(dead_code)]
+pub fn register_external_worker_with_card(ctx: &Arc<AppContext>, url: &str, model_card: ModelCard) {
+    let worker: Arc<dyn Worker> = Arc::new(
+        BasicWorkerBuilder::new(url)
+            .worker_type(WorkerType::Regular)
+            .runtime_type(RuntimeType::External)
+            .model(model_card)
+            .build(),
+    );
+
+    ctx.worker_registry.register(worker);
+}