server.rs 27.5 KB
Newer Older
1
use crate::{
2
    config::{ConnectionMode, HistoryBackend, RouterConfig},
3
    core::{WorkerRegistry, WorkerType},
4
    data_connector::{MemoryResponseStorage, NoOpResponseStorage, SharedResponseStorage},
5
6
7
8
9
10
11
12
13
14
15
16
    logging::{self, LoggingConfig},
    metrics::{self, PrometheusConfig},
    middleware::{self, QueuedRequest, TokenBucket},
    policies::PolicyRegistry,
    protocols::{
        spec::{
            ChatCompletionRequest, CompletionRequest, EmbeddingRequest, GenerateRequest,
            RerankRequest, ResponsesRequest, V1RerankReqInput,
        },
        worker_spec::{WorkerApiResponse, WorkerConfigRequest, WorkerErrorResponse},
    },
    reasoning_parser::ParserFactory,
17
    routers::WorkerInitializer,
18
19
20
21
22
23
24
    routers::{
        router_manager::{RouterId, RouterManager},
        RouterFactory, RouterTrait,
    },
    service_discovery::{start_service_discovery, ServiceDiscoveryConfig},
    tokenizer::{factory as tokenizer_factory, traits::Tokenizer},
    tool_parser::ParserRegistry,
25
};
26
use axum::{
27
    extract::{Path, Query, Request, State},
28
29
    http::StatusCode,
    response::{IntoResponse, Response},
30
    routing::{delete, get, post},
31
    serve, Json, Router,
32
};
33
use reqwest::Client;
34
35
36
37
38
39
40
41
use serde::Deserialize;
use serde_json::json;
use std::{
    sync::atomic::{AtomicBool, Ordering},
    sync::Arc,
    time::Duration,
};
use tokio::{net::TcpListener, signal, spawn};
42
use tracing::{error, info, warn, Level};
43

44
#[derive(Clone)]
45
pub struct AppContext {
46
    pub client: Client,
47
    pub router_config: RouterConfig,
48
    pub rate_limiter: Arc<TokenBucket>,
49
50
51
    pub tokenizer: Option<Arc<dyn Tokenizer>>,
    pub reasoning_parser_factory: Option<ParserFactory>,
    pub tool_parser_registry: Option<&'static ParserRegistry>,
52
53
54
    pub worker_registry: Arc<WorkerRegistry>,
    pub policy_registry: Arc<PolicyRegistry>,
    pub router_manager: Option<Arc<RouterManager>>,
55
    pub response_storage: SharedResponseStorage,
56
57
}

58
impl AppContext {
59
60
61
62
    pub fn new(
        router_config: RouterConfig,
        client: Client,
        max_concurrent_requests: usize,
63
        rate_limit_tokens_per_second: Option<usize>,
64
    ) -> Result<Self, String> {
65
66
        let rate_limit_tokens = rate_limit_tokens_per_second.unwrap_or(max_concurrent_requests);
        let rate_limiter = Arc::new(TokenBucket::new(max_concurrent_requests, rate_limit_tokens));
67
68
69

        // Initialize gRPC-specific components only when in gRPC mode
        let (tokenizer, reasoning_parser_factory, tool_parser_registry) =
70
            if router_config.connection_mode == ConnectionMode::Grpc {
71
72
73
74
75
76
77
78
79
80
81
82
83
                // Get tokenizer path (required for gRPC mode)
                let tokenizer_path = router_config
                    .tokenizer_path
                    .clone()
                    .or_else(|| router_config.model_path.clone())
                    .ok_or_else(|| {
                        "gRPC mode requires either --tokenizer-path or --model-path to be specified"
                            .to_string()
                    })?;

                // Initialize all gRPC components
                let tokenizer = Some(
                    tokenizer_factory::create_tokenizer(&tokenizer_path)
84
                        .map_err(|e| format!("Failed to create tokenizer: {e}"))?,
85
86
87
88
89
90
91
92
93
94
                );
                let reasoning_parser_factory = Some(ParserFactory::new());
                let tool_parser_registry = Some(ParserRegistry::new());

                (tokenizer, reasoning_parser_factory, tool_parser_registry)
            } else {
                // HTTP mode doesn't need these components
                (None, None, None)
            };

95
        let worker_registry = Arc::new(WorkerRegistry::new());
96
        let policy_registry = Arc::new(PolicyRegistry::new(router_config.policy.clone()));
97

98
        let router_manager = None;
99

100
101
102
103
104
105
        // Initialize response storage based on configuration
        let response_storage: SharedResponseStorage = match router_config.history_backend {
            HistoryBackend::Memory => Arc::new(MemoryResponseStorage::new()),
            HistoryBackend::None => Arc::new(NoOpResponseStorage::new()),
        };

106
        Ok(Self {
107
            client,
108
            router_config,
109
            rate_limiter,
110
111
112
            tokenizer,
            reasoning_parser_factory,
            tool_parser_registry,
113
114
115
            worker_registry,
            policy_registry,
            router_manager,
116
            response_storage,
117
        })
118
119
120
    }
}

121
122
123
124
#[derive(Clone)]
pub struct AppState {
    pub router: Arc<dyn RouterTrait>,
    pub context: Arc<AppContext>,
125
    pub concurrency_queue_tx: Option<tokio::sync::mpsc::Sender<QueuedRequest>>,
126
    pub router_manager: Option<Arc<RouterManager>>,
127
128
}

129
130
131
// Fallback handler for unmatched routes
async fn sink_handler() -> Response {
    StatusCode::NOT_FOUND.into_response()
132
133
}

134
135
136
// Health check endpoints
async fn liveness(State(state): State<Arc<AppState>>) -> Response {
    state.router.liveness()
137
138
}

139
140
async fn readiness(State(state): State<Arc<AppState>>) -> Response {
    state.router.readiness()
141
142
}

143
async fn health(State(state): State<Arc<AppState>>, req: Request) -> Response {
144
    state.router.health(req).await
145
146
}

147
async fn health_generate(State(state): State<Arc<AppState>>, req: Request) -> Response {
148
    state.router.health_generate(req).await
149
150
}

151
async fn get_server_info(State(state): State<Arc<AppState>>, req: Request) -> Response {
152
    state.router.get_server_info(req).await
153
154
}

155
async fn v1_models(State(state): State<Arc<AppState>>, req: Request) -> Response {
156
    state.router.get_models(req).await
157
158
}

159
async fn get_model_info(State(state): State<Arc<AppState>>, req: Request) -> Response {
160
    state.router.get_model_info(req).await
161
}
162

163
164
// Generation endpoints
// The RouterTrait now accepts optional headers and typed body directly
165
async fn generate(
166
167
168
169
    State(state): State<Arc<AppState>>,
    headers: http::HeaderMap,
    Json(body): Json<GenerateRequest>,
) -> Response {
170
171
172
173
    state
        .router
        .route_generate(Some(&headers), &body, None)
        .await
174
175
176
}

async fn v1_chat_completions(
177
178
179
180
    State(state): State<Arc<AppState>>,
    headers: http::HeaderMap,
    Json(body): Json<ChatCompletionRequest>,
) -> Response {
181
    state.router.route_chat(Some(&headers), &body, None).await
182
183
184
}

async fn v1_completions(
185
186
187
188
    State(state): State<Arc<AppState>>,
    headers: http::HeaderMap,
    Json(body): Json<CompletionRequest>,
) -> Response {
189
190
191
192
    state
        .router
        .route_completion(Some(&headers), &body, None)
        .await
193
194
}

195
196
197
198
199
async fn rerank(
    State(state): State<Arc<AppState>>,
    headers: http::HeaderMap,
    Json(body): Json<RerankRequest>,
) -> Response {
200
    state.router.route_rerank(Some(&headers), &body, None).await
201
202
203
204
205
206
207
208
209
}

async fn v1_rerank(
    State(state): State<Arc<AppState>>,
    headers: http::HeaderMap,
    Json(body): Json<V1RerankReqInput>,
) -> Response {
    state
        .router
210
        .route_rerank(Some(&headers), &body.into(), None)
211
212
213
        .await
}

214
215
216
217
218
async fn v1_responses(
    State(state): State<Arc<AppState>>,
    headers: http::HeaderMap,
    Json(body): Json<ResponsesRequest>,
) -> Response {
219
220
221
222
    state
        .router
        .route_responses(Some(&headers), &body, None)
        .await
223
224
}

225
226
227
228
229
230
231
232
233
234
235
async fn v1_embeddings(
    State(state): State<Arc<AppState>>,
    headers: http::HeaderMap,
    Json(body): Json<EmbeddingRequest>,
) -> Response {
    state
        .router
        .route_embeddings(Some(&headers), &body, None)
        .await
}

236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
async fn v1_responses_get(
    State(state): State<Arc<AppState>>,
    Path(response_id): Path<String>,
    headers: http::HeaderMap,
) -> Response {
    state
        .router
        .get_response(Some(&headers), &response_id)
        .await
}

async fn v1_responses_cancel(
    State(state): State<Arc<AppState>>,
    Path(response_id): Path<String>,
    headers: http::HeaderMap,
) -> Response {
    state
        .router
        .cancel_response(Some(&headers), &response_id)
        .await
}

async fn v1_responses_delete(
    State(state): State<Arc<AppState>>,
    Path(response_id): Path<String>,
    headers: http::HeaderMap,
) -> Response {
    // Python server does not support this yet
    state
        .router
        .delete_response(Some(&headers), &response_id)
        .await
}

async fn v1_responses_list_input_items(
    State(state): State<Arc<AppState>>,
    Path(response_id): Path<String>,
    headers: http::HeaderMap,
) -> Response {
    // Python server does not support this yet
    state
        .router
        .list_response_input_items(Some(&headers), &response_id)
        .await
}

282
283
284
// ---------- Worker management endpoints (Legacy) ----------

#[derive(Deserialize)]
285
struct AddWorkerQuery {
286
    url: String,
287
    api_key: Option<String>,
288
289
}

290
async fn add_worker(
291
    State(state): State<Arc<AppState>>,
292
    Query(AddWorkerQuery { url, api_key }): Query<AddWorkerQuery>,
293
) -> Response {
294
    match state.router.add_worker(&url, &api_key).await {
295
296
        Ok(message) => (StatusCode::OK, message).into_response(),
        Err(error) => (StatusCode::BAD_REQUEST, error).into_response(),
297
    }
298
299
}

300
301
302
async fn list_workers(State(state): State<Arc<AppState>>) -> Response {
    let worker_list = state.router.get_worker_urls();
    Json(serde_json::json!({ "urls": worker_list })).into_response()
303
304
}

305
async fn remove_worker(
306
    State(state): State<Arc<AppState>>,
307
    Query(AddWorkerQuery { url, .. }): Query<AddWorkerQuery>,
308
) -> Response {
309
    state.router.remove_worker(&url);
310
311
    (
        StatusCode::OK,
312
        format!("Successfully removed worker: {url}"),
313
314
    )
        .into_response()
315
316
}

317
async fn flush_cache(State(state): State<Arc<AppState>>, _req: Request) -> Response {
318
    state.router.flush_cache().await
319
320
}

321
async fn get_loads(State(state): State<Arc<AppState>>, _req: Request) -> Response {
322
    state.router.get_worker_loads().await
323
324
}

325
// ---------- Worker management endpoints (RESTful) ----------
326
327
328
329
330
331

/// POST /workers - Add a new worker with full configuration
async fn create_worker(
    State(state): State<Arc<AppState>>,
    Json(config): Json<WorkerConfigRequest>,
) -> Response {
332
333
    // Check if we have a RouterManager (enable_igw=true)
    if let Some(router_manager) = &state.router_manager {
334
        // Call RouterManager's add_worker method directly with the full config
335
336
337
338
339
340
        match router_manager.add_worker(config).await {
            Ok(response) => (StatusCode::OK, Json(response)).into_response(),
            Err(error) => (StatusCode::BAD_REQUEST, Json(error)).into_response(),
        }
    } else {
        // In single router mode, use the router's add_worker with basic config
341
        match state.router.add_worker(&config.url, &config.api_key).await {
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
            Ok(message) => {
                let response = WorkerApiResponse {
                    success: true,
                    message,
                    worker: None,
                };
                (StatusCode::OK, Json(response)).into_response()
            }
            Err(error) => {
                let error_response = WorkerErrorResponse {
                    error,
                    code: "ADD_WORKER_FAILED".to_string(),
                };
                (StatusCode::BAD_REQUEST, Json(error_response)).into_response()
            }
        }
    }
}

/// GET /workers - List all workers with details
async fn list_workers_rest(State(state): State<Arc<AppState>>) -> Response {
363
    if let Some(router_manager) = &state.router_manager {
364
365
366
367
368
369
370
371
372
373
        let response = router_manager.list_workers();
        Json(response).into_response()
    } else {
        // In single router mode, get detailed worker info from registry
        let workers = state.context.worker_registry.get_all();
        let response = serde_json::json!({
            "workers": workers.iter().map(|worker| {
                let mut worker_info = serde_json::json!({
                    "url": worker.url(),
                    "model_id": worker.model_id(),
374
375
376
377
378
                    "worker_type": match worker.worker_type() {
                        WorkerType::Regular => "regular",
                        WorkerType::Prefill { .. } => "prefill",
                        WorkerType::Decode => "decode",
                    },
379
380
381
382
383
384
385
386
                    "is_healthy": worker.is_healthy(),
                    "load": worker.load(),
                    "connection_mode": format!("{:?}", worker.connection_mode()),
                    "priority": worker.priority(),
                    "cost": worker.cost(),
                });

                // Add bootstrap_port for Prefill workers
387
                if let WorkerType::Prefill { bootstrap_port } = worker.worker_type() {
388
389
390
391
392
393
394
395
396
                    worker_info["bootstrap_port"] = serde_json::json!(bootstrap_port);
                }

                worker_info
            }).collect::<Vec<_>>(),
            "total": workers.len(),
            "stats": {
                "prefill_count": state.context.worker_registry.get_prefill_workers().len(),
                "decode_count": state.context.worker_registry.get_decode_workers().len(),
397
                "regular_count": state.context.worker_registry.get_by_type(&WorkerType::Regular).len(),
398
399
400
401
402
403
404
            }
        });
        Json(response).into_response()
    }
}

/// GET /workers/{url} - Get specific worker info
405
async fn get_worker(State(state): State<Arc<AppState>>, Path(url): Path<String>) -> Response {
406
    if let Some(router_manager) = &state.router_manager {
407
408
409
410
        if let Some(worker) = router_manager.get_worker(&url) {
            Json(worker).into_response()
        } else {
            let error = WorkerErrorResponse {
411
                error: format!("Worker {url} not found"),
412
413
414
415
416
417
418
                code: "WORKER_NOT_FOUND".to_string(),
            };
            (StatusCode::NOT_FOUND, Json(error)).into_response()
        }
    } else {
        let workers = state.router.get_worker_urls();
        if workers.contains(&url) {
419
            Json(json!({
420
421
422
                "url": url,
                "model_id": "unknown",
                "is_healthy": true
423
424
            }))
            .into_response()
425
426
        } else {
            let error = WorkerErrorResponse {
427
                error: format!("Worker {url} not found"),
428
429
430
431
432
433
434
435
                code: "WORKER_NOT_FOUND".to_string(),
            };
            (StatusCode::NOT_FOUND, Json(error)).into_response()
        }
    }
}

/// DELETE /workers/{url} - Remove a worker
436
async fn delete_worker(State(state): State<Arc<AppState>>, Path(url): Path<String>) -> Response {
437
    if let Some(router_manager) = &state.router_manager {
438
439
440
441
442
443
444
445
446
        match router_manager.remove_worker_from_registry(&url) {
            Ok(response) => (StatusCode::OK, Json(response)).into_response(),
            Err(error) => (StatusCode::BAD_REQUEST, Json(error)).into_response(),
        }
    } else {
        // In single router mode, use router's remove_worker
        state.router.remove_worker(&url);
        let response = WorkerApiResponse {
            success: true,
447
            message: format!("Worker {url} removed successfully"),
448
449
450
451
452
453
            worker: None,
        };
        (StatusCode::OK, Json(response)).into_response()
    }
}

454
455
456
pub struct ServerConfig {
    pub host: String,
    pub port: u16,
457
    pub router_config: RouterConfig,
458
    pub max_payload_size: usize,
459
    pub log_dir: Option<String>,
460
    pub log_level: Option<String>,
461
    pub service_discovery_config: Option<ServiceDiscoveryConfig>,
462
    pub prometheus_config: Option<PrometheusConfig>,
463
    pub request_timeout_secs: u64,
464
    pub request_id_headers: Option<Vec<String>>,
465
466
}

467
468
469
470
471
472
473
474
475
476
477
/// Build the Axum application with all routes and middleware
pub fn build_app(
    app_state: Arc<AppState>,
    max_payload_size: usize,
    request_id_headers: Vec<String>,
    cors_allowed_origins: Vec<String>,
) -> Router {
    // Create routes
    let protected_routes = Router::new()
        .route("/generate", post(generate))
        .route("/v1/chat/completions", post(v1_chat_completions))
478
        .route("/v1/completions", post(v1_completions))
479
480
        .route("/rerank", post(rerank))
        .route("/v1/rerank", post(v1_rerank))
481
        .route("/v1/responses", post(v1_responses))
482
        .route("/v1/embeddings", post(v1_embeddings))
483
484
485
486
487
488
489
490
491
492
        .route("/v1/responses/{response_id}", get(v1_responses_get))
        .route(
            "/v1/responses/{response_id}/cancel",
            post(v1_responses_cancel),
        )
        .route("/v1/responses/{response_id}", delete(v1_responses_delete))
        .route(
            "/v1/responses/{response_id}/input",
            get(v1_responses_list_input_items),
        )
493
494
        .route_layer(axum::middleware::from_fn_with_state(
            app_state.clone(),
495
            middleware::concurrency_limit_middleware,
496
        ));
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513

    let public_routes = Router::new()
        .route("/liveness", get(liveness))
        .route("/readiness", get(readiness))
        .route("/health", get(health))
        .route("/health_generate", get(health_generate))
        .route("/v1/models", get(v1_models))
        .route("/get_model_info", get(get_model_info))
        .route("/get_server_info", get(get_server_info));

    let admin_routes = Router::new()
        .route("/add_worker", post(add_worker))
        .route("/remove_worker", post(remove_worker))
        .route("/list_workers", get(list_workers))
        .route("/flush_cache", post(flush_cache))
        .route("/get_loads", get(get_loads));

514
515
516
517
518
    // Worker management routes
    let worker_routes = Router::new()
        .route("/workers", post(create_worker))
        .route("/workers", get(list_workers_rest))
        .route("/workers/{url}", get(get_worker))
519
        .route("/workers/{url}", delete(delete_worker));
520

521
522
523
524
525
    // Build app with all routes and middleware
    Router::new()
        .merge(protected_routes)
        .merge(public_routes)
        .merge(admin_routes)
526
        .merge(worker_routes)
527
528
529
530
        // Request body size limiting
        .layer(tower_http::limit::RequestBodyLimitLayer::new(
            max_payload_size,
        ))
531
532
        .layer(middleware::create_logging_layer())
        .layer(middleware::RequestIdLayer::new(request_id_headers))
533
534
535
536
537
538
        .layer(create_cors_layer(cors_allowed_origins))
        .fallback(sink_handler)
        .with_state(app_state)
}

pub async fn startup(config: ServerConfig) -> Result<(), Box<dyn std::error::Error>> {
539
540
541
542
543
    // Only initialize logging if not already done (for Python bindings support)
    static LOGGING_INITIALIZED: AtomicBool = AtomicBool::new(false);

    let _log_guard = if !LOGGING_INITIALIZED.swap(true, Ordering::SeqCst) {
        Some(logging::init_logging(LoggingConfig {
544
545
546
547
548
549
            level: config
                .log_level
                .as_deref()
                .and_then(|s| match s.to_uppercase().parse::<Level>() {
                    Ok(l) => Some(l),
                    Err(_) => {
550
                        warn!("Invalid log level string: '{s}'. Defaulting to INFO.");
551
552
553
554
                        None
                    }
                })
                .unwrap_or(Level::INFO),
555
556
557
558
559
560
561
562
563
            json_format: false,
            log_dir: config.log_dir.clone(),
            colorize: true,
            log_file_name: "sgl-router".to_string(),
            log_targets: None,
        }))
    } else {
        None
    };
564

565
566
    // Initialize prometheus metrics exporter
    if let Some(prometheus_config) = config.prometheus_config {
567
        metrics::start_prometheus(prometheus_config);
568
569
    }

570
    info!(
571
572
573
574
575
        "Starting router on {}:{} | mode: {:?} | policy: {:?} | max_payload: {}MB",
        config.host,
        config.port,
        config.router_config.mode,
        config.router_config.policy,
576
577
578
        config.max_payload_size / (1024 * 1024)
    );

579
    let client = Client::builder()
580
        .pool_idle_timeout(Some(Duration::from_secs(50)))
581
        .pool_max_idle_per_host(500)
582
        .timeout(Duration::from_secs(config.request_timeout_secs))
583
        .connect_timeout(Duration::from_secs(10))
584
        .tcp_nodelay(true)
585
        .tcp_keepalive(Some(Duration::from_secs(30)))
586
587
588
        .build()
        .expect("Failed to create HTTP client");

589
    // Create the application context with all dependencies
590
    let app_context = AppContext::new(
591
592
593
        config.router_config.clone(),
        client.clone(),
        config.router_config.max_concurrent_requests,
594
        config.router_config.rate_limit_tokens_per_second,
595
596
597
598
    )?;

    let app_context = Arc::new(app_context);

599
600
601
602
    info!(
        "Initializing workers for routing mode: {:?}",
        config.router_config.mode
    );
603
604
605
606
607
608
609
    WorkerInitializer::initialize_workers(
        &config.router_config,
        &app_context.worker_registry,
        Some(&app_context.policy_registry),
    )
    .await
    .map_err(|e| format!("Failed to initialize workers: {}", e))?;
610
611
612
613
614
615
616

    let worker_stats = app_context.worker_registry.stats();
    info!(
        "Workers initialized: {} total, {} healthy",
        worker_stats.total_workers, worker_stats.healthy_workers
    );

617
    // Create the appropriate router based on enable_igw flag
618
619
620
621
622
623
624
625
626
627
628
629
630
    let (router, router_manager): (Arc<dyn RouterTrait>, Option<Arc<RouterManager>>) =
        if config.router_config.enable_igw {
            info!("Multi-router mode enabled (enable_igw=true)");

            // Create RouterManager with shared registries from AppContext
            let router_manager = Arc::new(RouterManager::new(
                config.router_config.clone(),
                client.clone(),
                app_context.worker_registry.clone(),
                app_context.policy_registry.clone(),
            ));

            // 1. HTTP Regular Router
631
            match RouterFactory::create_regular_router(&app_context).await {
632
633
634
635
636
637
638
639
640
641
                Ok(http_regular) => {
                    info!("Created HTTP Regular router");
                    router_manager.register_router(
                        RouterId::new("http-regular".to_string()),
                        Arc::from(http_regular),
                    );
                }
                Err(e) => {
                    warn!("Failed to create HTTP Regular router: {e}");
                }
642
643
            }

644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
            // 2. HTTP PD Router
            match RouterFactory::create_pd_router(
                None,
                None,
                &config.router_config.policy,
                &app_context,
            )
            .await
            {
                Ok(http_pd) => {
                    info!("Created HTTP PD router");
                    router_manager
                        .register_router(RouterId::new("http-pd".to_string()), Arc::from(http_pd));
                }
                Err(e) => {
                    warn!("Failed to create HTTP PD router: {e}");
                }
661
            }
662

663
            // TODO: Add gRPC routers once we have dynamic tokenizer loading
664

665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
            info!(
                "RouterManager initialized with {} routers",
                router_manager.router_count()
            );
            (
                router_manager.clone() as Arc<dyn RouterTrait>,
                Some(router_manager),
            )
        } else {
            info!("Single router mode (enable_igw=false)");
            // Create single router with the context
            (
                Arc::from(RouterFactory::create_router(&app_context).await?),
                None,
            )
        };
681
682
683
684
685
686
687
688
689

    // Start health checker for all workers in the registry
    let _health_checker = app_context
        .worker_registry
        .start_health_checker(config.router_config.health_check.check_interval_secs);
    info!(
        "Started health checker for workers with {}s interval",
        config.router_config.health_check.check_interval_secs
    );
690

691
    // Set up concurrency limiter with queue if configured
692
    let (limiter, processor) = middleware::ConcurrencyLimiter::new(
693
694
695
696
697
698
699
        app_context.rate_limiter.clone(),
        config.router_config.queue_size,
        Duration::from_secs(config.router_config.queue_timeout_secs),
    );

    // Start queue processor if enabled
    if let Some(processor) = processor {
700
        spawn(processor.run());
701
702
703
704
705
706
        info!(
            "Started request queue with size: {}, timeout: {}s",
            config.router_config.queue_size, config.router_config.queue_timeout_secs
        );
    }

707
708
    // Create app state with router and context
    let app_state = Arc::new(AppState {
709
        router,
710
        context: app_context.clone(),
711
        concurrency_queue_tx: limiter.queue_tx.clone(),
712
        router_manager,
713
    });
714
    let router_arc = Arc::clone(&app_state.router);
715

716
717
718
    // Start the service discovery if enabled
    if let Some(service_discovery_config) = config.service_discovery_config {
        if service_discovery_config.enabled {
719
            match start_service_discovery(service_discovery_config, router_arc).await {
720
                Ok(handle) => {
721
                    info!("Service discovery started");
722
723
724
725
726
727
728
729
                    // Spawn a task to handle the service discovery thread
                    spawn(async move {
                        if let Err(e) = handle.await {
                            error!("Service discovery task failed: {:?}", e);
                        }
                    });
                }
                Err(e) => {
730
                    error!("Failed to start service discovery: {e}");
731
732
733
734
735
736
                    warn!("Continuing without service discovery");
                }
            }
        }
    }

737
    info!(
738
        "Router ready | workers: {:?}",
739
        app_state.router.get_worker_urls()
740
    );
741

742
743
744
745
746
747
748
749
750
    let request_id_headers = 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(),
        ]
    });

751
752
753
754
755
756
757
    // Build the application
    let app = build_app(
        app_state,
        config.max_payload_size,
        request_id_headers,
        config.router_config.cors_allowed_origins.clone(),
    );
758

759
760
761
    let addr = format!("{}:{}", config.host, config.port);
    let listener = TcpListener::bind(&addr).await?;
    info!("Starting server on {}", addr);
762
    serve(listener, app)
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
        .with_graceful_shutdown(shutdown_signal())
        .await
        .map_err(|e| Box::new(e) as Box<dyn std::error::Error>)?;

    Ok(())
}

// Graceful shutdown handler
async fn shutdown_signal() {
    let ctrl_c = async {
        signal::ctrl_c()
            .await
            .expect("failed to install Ctrl+C handler");
    };

    #[cfg(unix)]
    let terminate = async {
        signal::unix::signal(signal::unix::SignalKind::terminate())
            .expect("failed to install signal handler")
            .recv()
            .await;
    };

    #[cfg(not(unix))]
    let terminate = std::future::pending::<()>();

    tokio::select! {
        _ = ctrl_c => {
            info!("Received Ctrl+C, starting graceful shutdown");
        },
        _ = terminate => {
            info!("Received terminate signal, starting graceful shutdown");
        },
    }
}

// CORS Layer Creation
fn create_cors_layer(allowed_origins: Vec<String>) -> tower_http::cors::CorsLayer {
    use tower_http::cors::Any;

    let cors = if allowed_origins.is_empty() {
        // Allow all origins if none specified
        tower_http::cors::CorsLayer::new()
            .allow_origin(Any)
            .allow_methods(Any)
            .allow_headers(Any)
            .expose_headers(Any)
    } else {
        // Restrict to specific origins
        let origins: Vec<http::HeaderValue> = allowed_origins
            .into_iter()
            .filter_map(|origin| origin.parse().ok())
            .collect();

        tower_http::cors::CorsLayer::new()
            .allow_origin(origins)
            .allow_methods([http::Method::GET, http::Method::POST, http::Method::OPTIONS])
            .allow_headers([http::header::CONTENT_TYPE, http::header::AUTHORIZATION])
            .expose_headers([http::header::HeaderName::from_static("x-request-id")])
    };

    cors.max_age(Duration::from_secs(3600))
825
}