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

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

56
impl AppContext {
57
58
59
60
    pub fn new(
        router_config: RouterConfig,
        client: Client,
        max_concurrent_requests: usize,
61
        rate_limit_tokens_per_second: Option<usize>,
62
    ) -> Result<Self, String> {
63
64
        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));
65
66

        let (tokenizer, reasoning_parser_factory, tool_parser_registry) =
67
            if router_config.connection_mode == ConnectionMode::Grpc {
68
69
70
71
72
73
74
75
76
77
78
                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()
                    })?;

                let tokenizer = Some(
                    tokenizer_factory::create_tokenizer(&tokenizer_path)
79
                        .map_err(|e| format!("Failed to create tokenizer: {e}"))?,
80
81
82
83
84
85
86
87
88
                );
                let reasoning_parser_factory = Some(ParserFactory::new());
                let tool_parser_registry = Some(ParserRegistry::new());

                (tokenizer, reasoning_parser_factory, tool_parser_registry)
            } else {
                (None, None, None)
            };

89
        let worker_registry = Arc::new(WorkerRegistry::new());
90
        let policy_registry = Arc::new(PolicyRegistry::new(router_config.policy.clone()));
91

92
        let router_manager = None;
93

94
95
96
        let response_storage: SharedResponseStorage = match router_config.history_backend {
            HistoryBackend::Memory => Arc::new(MemoryResponseStorage::new()),
            HistoryBackend::None => Arc::new(NoOpResponseStorage::new()),
97
98
99
100
101
102
103
104
105
106
107
            HistoryBackend::Oracle => {
                let oracle_cfg = router_config.oracle.clone().ok_or_else(|| {
                    "oracle configuration is required when history_backend=oracle".to_string()
                })?;

                let storage = OracleResponseStorage::new(oracle_cfg).map_err(|err| {
                    format!("failed to initialize Oracle response storage: {err}")
                })?;

                Arc::new(storage)
            }
108
109
        };

110
        Ok(Self {
111
            client,
112
            router_config,
113
            rate_limiter,
114
115
116
            tokenizer,
            reasoning_parser_factory,
            tool_parser_registry,
117
118
119
            worker_registry,
            policy_registry,
            router_manager,
120
            response_storage,
121
        })
122
123
124
    }
}

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

133
134
async fn sink_handler() -> Response {
    StatusCode::NOT_FOUND.into_response()
135
136
}

137
138
async fn liveness() -> Response {
    (StatusCode::OK, "OK").into_response()
139
140
}

141
async fn readiness(State(state): State<Arc<AppState>>) -> Response {
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
    let workers = state.context.worker_registry.get_all();
    let healthy_workers: Vec<_> = workers.iter().filter(|w| w.is_healthy()).collect();

    let is_ready = if state.context.router_config.enable_igw {
        !healthy_workers.is_empty()
    } else {
        match &state.context.router_config.mode {
            RoutingMode::PrefillDecode { .. } => {
                let has_prefill = healthy_workers
                    .iter()
                    .any(|w| matches!(w.worker_type(), WorkerType::Prefill { .. }));
                let has_decode = healthy_workers
                    .iter()
                    .any(|w| matches!(w.worker_type(), WorkerType::Decode));
                has_prefill && has_decode
            }
            RoutingMode::Regular { .. } => !healthy_workers.is_empty(),
            RoutingMode::OpenAI { .. } => !healthy_workers.is_empty(),
        }
    };

    if is_ready {
        (
            StatusCode::OK,
            Json(json!({
                "status": "ready",
                "healthy_workers": healthy_workers.len(),
                "total_workers": workers.len()
            })),
        )
            .into_response()
    } else {
        (
            StatusCode::SERVICE_UNAVAILABLE,
            Json(json!({
                "status": "not ready",
                "reason": "insufficient healthy workers"
            })),
        )
            .into_response()
    }
183
184
}

185
186
async fn health(_state: State<Arc<AppState>>) -> Response {
    liveness().await
187
188
}

189
async fn health_generate(State(state): State<Arc<AppState>>, req: Request) -> Response {
190
    state.router.health_generate(req).await
191
192
}

193
async fn get_server_info(State(state): State<Arc<AppState>>, req: Request) -> Response {
194
    state.router.get_server_info(req).await
195
196
}

197
async fn v1_models(State(state): State<Arc<AppState>>, req: Request) -> Response {
198
    state.router.get_models(req).await
199
200
}

201
async fn get_model_info(State(state): State<Arc<AppState>>, req: Request) -> Response {
202
    state.router.get_model_info(req).await
203
}
204

205
async fn generate(
206
207
208
209
    State(state): State<Arc<AppState>>,
    headers: http::HeaderMap,
    Json(body): Json<GenerateRequest>,
) -> Response {
210
211
212
213
    state
        .router
        .route_generate(Some(&headers), &body, None)
        .await
214
215
216
}

async fn v1_chat_completions(
217
218
219
220
    State(state): State<Arc<AppState>>,
    headers: http::HeaderMap,
    Json(body): Json<ChatCompletionRequest>,
) -> Response {
221
    state.router.route_chat(Some(&headers), &body, None).await
222
223
224
}

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

235
236
237
238
239
async fn rerank(
    State(state): State<Arc<AppState>>,
    headers: http::HeaderMap,
    Json(body): Json<RerankRequest>,
) -> Response {
240
    state.router.route_rerank(Some(&headers), &body, None).await
241
242
243
244
245
246
247
248
249
}

async fn v1_rerank(
    State(state): State<Arc<AppState>>,
    headers: http::HeaderMap,
    Json(body): Json<V1RerankReqInput>,
) -> Response {
    state
        .router
250
        .route_rerank(Some(&headers), &body.into(), None)
251
252
253
        .await
}

254
255
256
257
258
async fn v1_responses(
    State(state): State<Arc<AppState>>,
    headers: http::HeaderMap,
    Json(body): Json<ResponsesRequest>,
) -> Response {
259
260
261
262
    state
        .router
        .route_responses(Some(&headers), &body, None)
        .await
263
264
}

265
266
267
268
269
270
271
272
273
274
275
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
}

276
277
278
279
async fn v1_responses_get(
    State(state): State<Arc<AppState>>,
    Path(response_id): Path<String>,
    headers: http::HeaderMap,
280
    Query(params): Query<ResponsesGetParams>,
281
282
283
) -> Response {
    state
        .router
284
        .get_response(Some(&headers), &response_id, &params)
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
        .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 {
    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 {
    state
        .router
        .list_response_input_items(Some(&headers), &response_id)
        .await
}

321
#[derive(Deserialize)]
322
struct AddWorkerQuery {
323
    url: String,
324
    api_key: Option<String>,
325
326
}

327
async fn add_worker(
328
    State(state): State<Arc<AppState>>,
329
    Query(AddWorkerQuery { url, api_key }): Query<AddWorkerQuery>,
330
) -> Response {
331
332
333
334
335
336
337
338
339
340
    // Warn if router has API key but worker is being added without one
    if state.context.router_config.api_key.is_some() && api_key.is_none() {
        warn!(
            "Adding worker {} without API key while router has API key configured. \
            Worker will be accessible without authentication. \
            If the worker requires the same API key as the router, please specify it explicitly.",
            url
        );
    }

341
342
343
    let result = WorkerManager::add_worker(&url, &api_key, &state.context).await;

    match result {
344
345
        Ok(message) => (StatusCode::OK, message).into_response(),
        Err(error) => (StatusCode::BAD_REQUEST, error).into_response(),
346
    }
347
348
}

349
async fn list_workers(State(state): State<Arc<AppState>>) -> Response {
350
351
    let worker_list = WorkerManager::get_worker_urls(&state.context.worker_registry);
    Json(json!({ "urls": worker_list })).into_response()
352
353
}

354
async fn remove_worker(
355
    State(state): State<Arc<AppState>>,
356
    Query(AddWorkerQuery { url, .. }): Query<AddWorkerQuery>,
357
) -> Response {
358
359
360
361
362
363
    let result = WorkerManager::remove_worker(&url, &state.context);

    match result {
        Ok(message) => (StatusCode::OK, message).into_response(),
        Err(error) => (StatusCode::BAD_REQUEST, error).into_response(),
    }
364
365
}

366
async fn flush_cache(State(state): State<Arc<AppState>>, _req: Request) -> Response {
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
    match WorkerManager::flush_cache_all(&state.context.worker_registry, &state.context.client)
        .await
    {
        Ok(result) => {
            if result.failed.is_empty() {
                (
                    StatusCode::OK,
                    Json(json!({
                        "status": "success",
                        "message": result.message,
                        "workers_flushed": result.successful.len(),
                        "total_http_workers": result.http_workers,
                        "total_workers": result.total_workers
                    })),
                )
                    .into_response()
            } else {
                (
                    StatusCode::PARTIAL_CONTENT,
                    Json(json!({
                        "status": "partial_success",
                        "message": result.message,
                        "successful": result.successful,
                        "failed": result.failed.into_iter().map(|(url, err)| json!({
                            "worker": url,
                            "error": err
                        })).collect::<Vec<_>>(),
                        "total_http_workers": result.http_workers,
                        "total_workers": result.total_workers
                    })),
                )
                    .into_response()
            }
        }
        Err(e) => {
            error!("Failed to flush cache: {}", e);
            (
                StatusCode::INTERNAL_SERVER_ERROR,
                Json(json!({
                    "status": "error",
                    "message": format!("Failed to flush cache: {}", e)
                })),
            )
                .into_response()
        }
    }
413
414
}

415
async fn get_loads(State(state): State<Arc<AppState>>, _req: Request) -> Response {
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
    let result =
        WorkerManager::get_all_worker_loads(&state.context.worker_registry, &state.context.client)
            .await;

    let loads: Vec<Value> = result
        .loads
        .iter()
        .map(|info| {
            json!({
                "worker": &info.worker,
                "load": info.load
            })
        })
        .collect();

    (
        StatusCode::OK,
        Json(json!({
            "workers": loads
        })),
    )
        .into_response()
438
439
}

440
441
442
443
async fn create_worker(
    State(state): State<Arc<AppState>>,
    Json(config): Json<WorkerConfigRequest>,
) -> Response {
444
445
446
447
448
449
450
451
452
453
    // Warn if router has API key but worker is being added without one
    if state.context.router_config.api_key.is_some() && config.api_key.is_none() {
        warn!(
            "Adding worker {} without API key while router has API key configured. \
            Worker will be accessible without authentication. \
            If the worker requires the same API key as the router, please specify it explicitly.",
            config.url
        );
    }

454
455
456
457
458
459
460
461
462
463
    let result = WorkerManager::add_worker_from_config(&config, &state.context).await;

    match result {
        Ok(message) => {
            let response = WorkerApiResponse {
                success: true,
                message,
                worker: None,
            };
            (StatusCode::OK, Json(response)).into_response()
464
        }
465
466
467
468
469
470
        Err(error) => {
            let error_response = WorkerErrorResponse {
                error,
                code: "ADD_WORKER_FAILED".to_string(),
            };
            (StatusCode::BAD_REQUEST, Json(error_response)).into_response()
471
472
473
474
475
        }
    }
}

async fn list_workers_rest(State(state): State<Arc<AppState>>) -> Response {
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
    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(),
                "worker_type": match worker.worker_type() {
                    WorkerType::Regular => "regular",
                    WorkerType::Prefill { .. } => "prefill",
                    WorkerType::Decode => "decode",
                },
                "is_healthy": worker.is_healthy(),
                "load": worker.load(),
                "connection_mode": format!("{:?}", worker.connection_mode()),
                "priority": worker.priority(),
                "cost": worker.cost(),
            });

            if let WorkerType::Prefill { bootstrap_port } = worker.worker_type() {
                worker_info["bootstrap_port"] = serde_json::json!(bootstrap_port);
496
            }
497
498
499
500
501
502
503
504
505
506
507

            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(),
            "regular_count": state.context.worker_registry.get_by_type(&WorkerType::Regular).len(),
        }
    });
    Json(response).into_response()
508
509
}

510
async fn get_worker(State(state): State<Arc<AppState>>, Path(url): Path<String>) -> Response {
511
512
513
514
515
516
517
518
    let workers = WorkerManager::get_worker_urls(&state.context.worker_registry);
    if workers.contains(&url) {
        Json(json!({
            "url": url,
            "model_id": "unknown",
            "is_healthy": true
        }))
        .into_response()
519
    } else {
520
521
522
523
524
        let error = WorkerErrorResponse {
            error: format!("Worker {url} not found"),
            code: "WORKER_NOT_FOUND".to_string(),
        };
        (StatusCode::NOT_FOUND, Json(error)).into_response()
525
526
527
    }
}

528
async fn delete_worker(State(state): State<Arc<AppState>>, Path(url): Path<String>) -> Response {
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
    let result = WorkerManager::remove_worker(&url, &state.context);

    match result {
        Ok(message) => {
            let response = WorkerApiResponse {
                success: true,
                message,
                worker: None,
            };
            (StatusCode::OK, Json(response)).into_response()
        }
        Err(error) => {
            let error_response = WorkerErrorResponse {
                error,
                code: "REMOVE_WORKER_FAILED".to_string(),
            };
            (StatusCode::BAD_REQUEST, Json(error_response)).into_response()
546
547
548
549
        }
    }
}

550
551
552
pub struct ServerConfig {
    pub host: String,
    pub port: u16,
553
    pub router_config: RouterConfig,
554
    pub max_payload_size: usize,
555
    pub log_dir: Option<String>,
556
    pub log_level: Option<String>,
557
    pub service_discovery_config: Option<ServiceDiscoveryConfig>,
558
    pub prometheus_config: Option<PrometheusConfig>,
559
    pub request_timeout_secs: u64,
560
    pub request_id_headers: Option<Vec<String>>,
561
562
}

563
564
pub fn build_app(
    app_state: Arc<AppState>,
565
    auth_config: AuthConfig,
566
567
568
569
570
571
572
    max_payload_size: usize,
    request_id_headers: Vec<String>,
    cors_allowed_origins: Vec<String>,
) -> Router {
    let protected_routes = Router::new()
        .route("/generate", post(generate))
        .route("/v1/chat/completions", post(v1_chat_completions))
573
        .route("/v1/completions", post(v1_completions))
574
575
        .route("/rerank", post(rerank))
        .route("/v1/rerank", post(v1_rerank))
576
        .route("/v1/responses", post(v1_responses))
577
        .route("/v1/embeddings", post(v1_embeddings))
578
579
580
581
582
583
584
585
586
587
        .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),
        )
588
589
        .route_layer(axum::middleware::from_fn_with_state(
            app_state.clone(),
590
            middleware::concurrency_limit_middleware,
591
592
593
594
        ))
        .route_layer(axum::middleware::from_fn_with_state(
            auth_config.clone(),
            middleware::auth_middleware,
595
        ));
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610

    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))
611
612
613
614
615
        .route("/get_loads", get(get_loads))
        .route_layer(axum::middleware::from_fn_with_state(
            auth_config.clone(),
            middleware::auth_middleware,
        ));
616

617
618
619
620
    let worker_routes = Router::new()
        .route("/workers", post(create_worker))
        .route("/workers", get(list_workers_rest))
        .route("/workers/{url}", get(get_worker))
621
622
623
624
625
        .route("/workers/{url}", delete(delete_worker))
        .route_layer(axum::middleware::from_fn_with_state(
            auth_config.clone(),
            middleware::auth_middleware,
        ));
626

627
628
629
630
    Router::new()
        .merge(protected_routes)
        .merge(public_routes)
        .merge(admin_routes)
631
        .merge(worker_routes)
632
        .layer(axum::extract::DefaultBodyLimit::max(max_payload_size))
633
634
635
        .layer(tower_http::limit::RequestBodyLimitLayer::new(
            max_payload_size,
        ))
636
637
        .layer(middleware::create_logging_layer())
        .layer(middleware::RequestIdLayer::new(request_id_headers))
638
639
640
641
642
643
        .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>> {
644
645
646
647
    static LOGGING_INITIALIZED: AtomicBool = AtomicBool::new(false);

    let _log_guard = if !LOGGING_INITIALIZED.swap(true, Ordering::SeqCst) {
        Some(logging::init_logging(LoggingConfig {
648
649
650
651
652
653
            level: config
                .log_level
                .as_deref()
                .and_then(|s| match s.to_uppercase().parse::<Level>() {
                    Ok(l) => Some(l),
                    Err(_) => {
654
                        warn!("Invalid log level string: '{s}'. Defaulting to INFO.");
655
656
657
658
                        None
                    }
                })
                .unwrap_or(Level::INFO),
659
660
661
662
663
664
665
666
667
            json_format: false,
            log_dir: config.log_dir.clone(),
            colorize: true,
            log_file_name: "sgl-router".to_string(),
            log_targets: None,
        }))
    } else {
        None
    };
668

669
670
    if let Some(prometheus_config) = &config.prometheus_config {
        metrics::start_prometheus(prometheus_config.clone());
671
672
    }

673
    info!(
674
675
676
677
678
        "Starting router on {}:{} | mode: {:?} | policy: {:?} | max_payload: {}MB",
        config.host,
        config.port,
        config.router_config.mode,
        config.router_config.policy,
679
680
681
        config.max_payload_size / (1024 * 1024)
    );

682
    let client = Client::builder()
683
        .pool_idle_timeout(Some(Duration::from_secs(50)))
684
        .pool_max_idle_per_host(500)
685
        .timeout(Duration::from_secs(config.request_timeout_secs))
686
        .connect_timeout(Duration::from_secs(10))
687
        .tcp_nodelay(true)
688
        .tcp_keepalive(Some(Duration::from_secs(30)))
689
690
691
        .build()
        .expect("Failed to create HTTP client");

692
    let app_context = AppContext::new(
693
694
695
        config.router_config.clone(),
        client.clone(),
        config.router_config.max_concurrent_requests,
696
        config.router_config.rate_limit_tokens_per_second,
697
698
699
700
    )?;

    let app_context = Arc::new(app_context);

701
702
703
704
    info!(
        "Initializing workers for routing mode: {:?}",
        config.router_config.mode
    );
705
    WorkerManager::initialize_workers(
706
707
708
709
710
711
        &config.router_config,
        &app_context.worker_registry,
        Some(&app_context.policy_registry),
    )
    .await
    .map_err(|e| format!("Failed to initialize workers: {}", e))?;
712
713
714
715
716
717
718

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

719
720
    let router_manager = RouterManager::from_config(&config, &app_context).await?;
    let router: Arc<dyn RouterTrait> = router_manager.clone();
721
722
723
724
725
726
727
728

    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
    );
729

730
    let (limiter, processor) = middleware::ConcurrencyLimiter::new(
731
732
733
734
735
736
        app_context.rate_limiter.clone(),
        config.router_config.queue_size,
        Duration::from_secs(config.router_config.queue_timeout_secs),
    );

    if let Some(processor) = processor {
737
        spawn(processor.run());
738
739
740
741
742
743
        info!(
            "Started request queue with size: {}, timeout: {}s",
            config.router_config.queue_size, config.router_config.queue_timeout_secs
        );
    }

744
    let app_state = Arc::new(AppState {
745
        router,
746
        context: app_context.clone(),
747
        concurrency_queue_tx: limiter.queue_tx.clone(),
748
        router_manager: Some(router_manager),
749
    });
750
751
    if let Some(service_discovery_config) = config.service_discovery_config {
        if service_discovery_config.enabled {
752
753
            let app_context_arc = Arc::clone(&app_state.context);
            match start_service_discovery(service_discovery_config, app_context_arc).await {
754
                Ok(handle) => {
755
                    info!("Service discovery started");
756
757
758
759
760
761
762
                    spawn(async move {
                        if let Err(e) = handle.await {
                            error!("Service discovery task failed: {:?}", e);
                        }
                    });
                }
                Err(e) => {
763
                    error!("Failed to start service discovery: {e}");
764
765
766
767
768
769
                    warn!("Continuing without service discovery");
                }
            }
        }
    }

770
    info!(
771
        "Router ready | workers: {:?}",
772
        WorkerManager::get_worker_urls(&app_state.context.worker_registry)
773
    );
774

775
776
777
778
779
780
781
782
783
    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(),
        ]
    });

784
785
786
787
    let auth_config = AuthConfig {
        api_key: config.router_config.api_key.clone(),
    };

788
789
    let app = build_app(
        app_state,
790
        auth_config,
791
792
793
794
        config.max_payload_size,
        request_id_headers,
        config.router_config.cors_allowed_origins.clone(),
    );
795

796
797
798
    let addr = format!("{}:{}", config.host, config.port);
    let listener = TcpListener::bind(&addr).await?;
    info!("Starting server on {}", addr);
799
    serve(listener, app)
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
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
        .with_graceful_shutdown(shutdown_signal())
        .await
        .map_err(|e| Box::new(e) as Box<dyn std::error::Error>)?;

    Ok(())
}

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");
        },
    }
}

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

    let cors = if allowed_origins.is_empty() {
        tower_http::cors::CorsLayer::new()
            .allow_origin(Any)
            .allow_methods(Any)
            .allow_headers(Any)
            .expose_headers(Any)
    } else {
        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))
858
}