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

46
47
//

48
#[derive(Clone)]
49
pub struct AppContext {
50
    pub client: Client,
51
    pub router_config: RouterConfig,
52
    pub rate_limiter: Option<Arc<TokenBucket>>,
53
    pub tokenizer: Option<Arc<dyn Tokenizer>>,
54
    pub reasoning_parser_factory: Option<ReasoningParserFactory>,
55
    pub tool_parser_factory: Option<ToolParserFactory>,
56
57
58
    pub worker_registry: Arc<WorkerRegistry>,
    pub policy_registry: Arc<PolicyRegistry>,
    pub router_manager: Option<Arc<RouterManager>>,
59
    pub response_storage: SharedResponseStorage,
60
    pub conversation_storage: SharedConversationStorage,
61
    pub conversation_item_storage: crate::data_connector::SharedConversationItemStorage,
62
    pub load_monitor: Option<Arc<LoadMonitor>>,
63
64
    pub configured_reasoning_parser: Option<String>,
    pub configured_tool_parser: Option<String>,
65
66
}

67
impl AppContext {
68
69
70
    pub fn new(
        router_config: RouterConfig,
        client: Client,
71
72
        max_concurrent_requests: i32,
        rate_limit_tokens_per_second: Option<i32>,
73
    ) -> Result<Self, String> {
74
75
76
77
78
79
80
81
82
83
84
        let rate_limiter = match max_concurrent_requests {
            n if n <= 0 => None,
            n => {
                let rate_limit_tokens =
                    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,
                )))
            }
        };
85

86
87
88
89
90
91
92
93
94
95
96
97
        let (tokenizer, reasoning_parser_factory, tool_parser_factory) = if router_config
            .connection_mode
            == ConnectionMode::Grpc
        {
            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()
                })?;
98

99
100
101
102
103
104
105
106
107
108
109
110
111
            let tokenizer = Some(
                    tokenizer_factory::create_tokenizer_with_chat_template_blocking(
                        &tokenizer_path,
                        router_config.chat_template.as_deref(),
                    )
                    .map_err(|e| {
                        format!(
                            "Failed to create tokenizer from '{}': {}. \
                            Ensure the path is valid and points to a tokenizer file (tokenizer.json) \
                            or a HuggingFace model ID. For directories, ensure they contain tokenizer files.",
                            tokenizer_path, e
                        )
                    })?,
112
                );
113
114
            let reasoning_parser_factory = Some(crate::reasoning_parser::ParserFactory::new());
            let tool_parser_factory = Some(crate::tool_parser::ParserFactory::new());
115

116
117
118
119
            (tokenizer, reasoning_parser_factory, tool_parser_factory)
        } else {
            (None, None, None)
        };
120

121
        let worker_registry = Arc::new(WorkerRegistry::new());
122
        let policy_registry = Arc::new(PolicyRegistry::new(router_config.policy.clone()));
123

124
        let router_manager = None;
125

126
127
128
129
130
131
132
133
134
135
136
137
        let (response_storage, conversation_storage): (
            SharedResponseStorage,
            SharedConversationStorage,
        ) = match router_config.history_backend {
            HistoryBackend::Memory => (
                Arc::new(MemoryResponseStorage::new()),
                Arc::new(MemoryConversationStorage::new()),
            ),
            HistoryBackend::None => (
                Arc::new(NoOpResponseStorage::new()),
                Arc::new(NoOpConversationStorage::new()),
            ),
138
139
140
141
142
            HistoryBackend::Oracle => {
                let oracle_cfg = router_config.oracle.clone().ok_or_else(|| {
                    "oracle configuration is required when history_backend=oracle".to_string()
                })?;

143
144
145
146
147
                let response_storage =
                    OracleResponseStorage::new(oracle_cfg.clone()).map_err(|err| {
                        format!("failed to initialize Oracle response storage: {err}")
                    })?;

148
149
                let conversation_storage = OracleConversationStorage::new(oracle_cfg.clone())
                    .map_err(|err| {
150
151
                        format!("failed to initialize Oracle conversation storage: {err}")
                    })?;
152

153
                (Arc::new(response_storage), Arc::new(conversation_storage))
154
            }
155
156
        };

157
158
159
160
161
162
163
164
165
166
167
168
169
170
        // Conversation items storage (memory-backed for now)
        let conversation_item_storage: crate::data_connector::SharedConversationItemStorage =
            match router_config.history_backend {
                HistoryBackend::Oracle => {
                    let oracle_cfg = router_config.oracle.clone().ok_or_else(|| {
                        "oracle configuration is required when history_backend=oracle".to_string()
                    })?;
                    Arc::new(OracleConversationItemStorage::new(oracle_cfg).map_err(|e| {
                        format!("failed to initialize Oracle conversation item storage: {e}")
                    })?)
                }
                _ => Arc::new(MemoryConversationItemStorage::new()),
            };

171
172
173
174
175
176
177
        let load_monitor = Some(Arc::new(LoadMonitor::new(
            worker_registry.clone(),
            policy_registry.clone(),
            client.clone(),
            router_config.worker_startup_check_interval_secs,
        )));

178
179
180
        let configured_reasoning_parser = router_config.reasoning_parser.clone();
        let configured_tool_parser = router_config.tool_call_parser.clone();

181
        Ok(Self {
182
            client,
183
            router_config,
184
            rate_limiter,
185
186
            tokenizer,
            reasoning_parser_factory,
187
            tool_parser_factory,
188
189
190
            worker_registry,
            policy_registry,
            router_manager,
191
            response_storage,
192
            conversation_storage,
193
            conversation_item_storage,
194
            load_monitor,
195
196
            configured_reasoning_parser,
            configured_tool_parser,
197
        })
198
199
200
    }
}

201
202
203
204
#[derive(Clone)]
pub struct AppState {
    pub router: Arc<dyn RouterTrait>,
    pub context: Arc<AppContext>,
205
    pub concurrency_queue_tx: Option<tokio::sync::mpsc::Sender<QueuedRequest>>,
206
    pub router_manager: Option<Arc<RouterManager>>,
207
208
}

209
210
async fn sink_handler() -> Response {
    StatusCode::NOT_FOUND.into_response()
211
212
}

213
214
async fn liveness() -> Response {
    (StatusCode::OK, "OK").into_response()
215
216
}

217
async fn readiness(State(state): State<Arc<AppState>>) -> Response {
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
    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()
    }
259
260
}

261
262
async fn health(_state: State<Arc<AppState>>) -> Response {
    liveness().await
263
264
}

265
async fn health_generate(State(state): State<Arc<AppState>>, req: Request) -> Response {
266
    state.router.health_generate(req).await
267
268
}

269
async fn get_server_info(State(state): State<Arc<AppState>>, req: Request) -> Response {
270
    state.router.get_server_info(req).await
271
272
}

273
async fn v1_models(State(state): State<Arc<AppState>>, req: Request) -> Response {
274
    state.router.get_models(req).await
275
276
}

277
async fn get_model_info(State(state): State<Arc<AppState>>, req: Request) -> Response {
278
    state.router.get_model_info(req).await
279
}
280

281
async fn generate(
282
283
284
285
    State(state): State<Arc<AppState>>,
    headers: http::HeaderMap,
    Json(body): Json<GenerateRequest>,
) -> Response {
286
287
288
289
    state
        .router
        .route_generate(Some(&headers), &body, None)
        .await
290
291
292
}

async fn v1_chat_completions(
293
294
    State(state): State<Arc<AppState>>,
    headers: http::HeaderMap,
295
    ValidatedJson(body): ValidatedJson<ChatCompletionRequest>,
296
) -> Response {
297
    state.router.route_chat(Some(&headers), &body, None).await
298
299
300
}

async fn v1_completions(
301
302
303
304
    State(state): State<Arc<AppState>>,
    headers: http::HeaderMap,
    Json(body): Json<CompletionRequest>,
) -> Response {
305
306
307
308
    state
        .router
        .route_completion(Some(&headers), &body, None)
        .await
309
310
}

311
312
313
314
315
async fn rerank(
    State(state): State<Arc<AppState>>,
    headers: http::HeaderMap,
    Json(body): Json<RerankRequest>,
) -> Response {
316
    state.router.route_rerank(Some(&headers), &body, None).await
317
318
319
320
321
322
323
324
325
}

async fn v1_rerank(
    State(state): State<Arc<AppState>>,
    headers: http::HeaderMap,
    Json(body): Json<V1RerankReqInput>,
) -> Response {
    state
        .router
326
        .route_rerank(Some(&headers), &body.into(), None)
327
328
329
        .await
}

330
331
332
333
334
async fn v1_responses(
    State(state): State<Arc<AppState>>,
    headers: http::HeaderMap,
    Json(body): Json<ResponsesRequest>,
) -> Response {
335
336
337
338
    state
        .router
        .route_responses(Some(&headers), &body, None)
        .await
339
340
}

341
342
343
344
345
346
347
348
349
350
351
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
}

352
353
354
355
async fn v1_responses_get(
    State(state): State<Arc<AppState>>,
    Path(response_id): Path<String>,
    headers: http::HeaderMap,
356
    Query(params): Query<ResponsesGetParams>,
357
358
359
) -> Response {
    state
        .router
360
        .get_response(Some(&headers), &response_id, &params)
361
362
363
364
365
366
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
        .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
}

397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
async fn v1_conversations_create(
    State(state): State<Arc<AppState>>,
    headers: http::HeaderMap,
    Json(body): Json<Value>,
) -> Response {
    state
        .router
        .create_conversation(Some(&headers), &body)
        .await
}

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

async fn v1_conversations_update(
    State(state): State<Arc<AppState>>,
    Path(conversation_id): Path<String>,
    headers: http::HeaderMap,
    Json(body): Json<Value>,
) -> Response {
    state
        .router
        .update_conversation(Some(&headers), &conversation_id, &body)
        .await
}

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

442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
#[derive(Deserialize, Default)]
struct ListItemsQuery {
    limit: Option<usize>,
    order: Option<String>,
    after: Option<String>,
}

async fn v1_conversations_list_items(
    State(state): State<Arc<AppState>>,
    Path(conversation_id): Path<String>,
    Query(ListItemsQuery {
        limit,
        order,
        after,
    }): Query<ListItemsQuery>,
    headers: http::HeaderMap,
) -> Response {
    state
        .router
        .list_conversation_items(Some(&headers), &conversation_id, limit, order, after)
        .await
}

465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
#[derive(Deserialize, Default)]
struct GetItemQuery {
    /// Additional fields to include in response (not yet implemented)
    include: Option<Vec<String>>,
}

async fn v1_conversations_create_items(
    State(state): State<Arc<AppState>>,
    Path(conversation_id): Path<String>,
    headers: http::HeaderMap,
    Json(body): Json<Value>,
) -> Response {
    state
        .router
        .create_conversation_items(Some(&headers), &conversation_id, &body)
        .await
}

async fn v1_conversations_get_item(
    State(state): State<Arc<AppState>>,
    Path((conversation_id, item_id)): Path<(String, String)>,
    Query(query): Query<GetItemQuery>,
    headers: http::HeaderMap,
) -> Response {
    state
        .router
        .get_conversation_item(Some(&headers), &conversation_id, &item_id, query.include)
        .await
}

async fn v1_conversations_delete_item(
    State(state): State<Arc<AppState>>,
    Path((conversation_id, item_id)): Path<(String, String)>,
    headers: http::HeaderMap,
) -> Response {
    state
        .router
        .delete_conversation_item(Some(&headers), &conversation_id, &item_id)
        .await
}

506
#[derive(Deserialize)]
507
struct AddWorkerQuery {
508
    url: String,
509
    api_key: Option<String>,
510
511
}

512
async fn add_worker(
513
    State(state): State<Arc<AppState>>,
514
    Query(AddWorkerQuery { url, api_key }): Query<AddWorkerQuery>,
515
) -> Response {
516
517
518
519
520
521
522
523
524
525
    // 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
        );
    }

526
527
528
    let result = WorkerManager::add_worker(&url, &api_key, &state.context).await;

    match result {
529
530
        Ok(message) => (StatusCode::OK, message).into_response(),
        Err(error) => (StatusCode::BAD_REQUEST, error).into_response(),
531
    }
532
533
}

534
async fn list_workers(State(state): State<Arc<AppState>>) -> Response {
535
536
    let worker_list = WorkerManager::get_worker_urls(&state.context.worker_registry);
    Json(json!({ "urls": worker_list })).into_response()
537
538
}

539
async fn remove_worker(
540
    State(state): State<Arc<AppState>>,
541
    Query(AddWorkerQuery { url, .. }): Query<AddWorkerQuery>,
542
) -> Response {
543
544
545
546
547
548
    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(),
    }
549
550
}

551
async fn flush_cache(State(state): State<Arc<AppState>>, _req: Request) -> Response {
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
    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()
        }
    }
598
599
}

600
async fn get_loads(State(state): State<Arc<AppState>>, _req: Request) -> Response {
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
    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()
623
624
}

625
626
627
628
async fn create_worker(
    State(state): State<Arc<AppState>>,
    Json(config): Json<WorkerConfigRequest>,
) -> Response {
629
630
631
632
633
634
635
636
637
638
    // 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
        );
    }

639
640
641
642
643
644
645
646
647
648
    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()
649
        }
650
651
652
653
654
655
        Err(error) => {
            let error_response = WorkerErrorResponse {
                error,
                code: "ADD_WORKER_FAILED".to_string(),
            };
            (StatusCode::BAD_REQUEST, Json(error_response)).into_response()
656
657
658
659
660
        }
    }
}

async fn list_workers_rest(State(state): State<Arc<AppState>>) -> Response {
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
    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);
681
            }
682
683
684
685
686
687
688
689
690
691
692

            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()
693
694
}

695
async fn get_worker(State(state): State<Arc<AppState>>, Path(url): Path<String>) -> Response {
696
697
698
699
700
701
702
703
    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()
704
    } else {
705
706
707
708
709
        let error = WorkerErrorResponse {
            error: format!("Worker {url} not found"),
            code: "WORKER_NOT_FOUND".to_string(),
        };
        (StatusCode::NOT_FOUND, Json(error)).into_response()
710
711
712
    }
}

713
async fn delete_worker(State(state): State<Arc<AppState>>, Path(url): Path<String>) -> Response {
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
    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()
731
732
733
734
        }
    }
}

735
736
737
pub struct ServerConfig {
    pub host: String,
    pub port: u16,
738
    pub router_config: RouterConfig,
739
    pub max_payload_size: usize,
740
    pub log_dir: Option<String>,
741
    pub log_level: Option<String>,
742
    pub service_discovery_config: Option<ServiceDiscoveryConfig>,
743
    pub prometheus_config: Option<PrometheusConfig>,
744
    pub request_timeout_secs: u64,
745
    pub request_id_headers: Option<Vec<String>>,
746
747
}

748
749
pub fn build_app(
    app_state: Arc<AppState>,
750
    auth_config: AuthConfig,
751
752
753
754
755
756
757
    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))
758
        .route("/v1/completions", post(v1_completions))
759
760
        .route("/rerank", post(rerank))
        .route("/v1/rerank", post(v1_rerank))
761
        .route("/v1/responses", post(v1_responses))
762
        .route("/v1/embeddings", post(v1_embeddings))
763
764
765
766
767
768
769
770
771
772
        .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),
        )
773
774
775
776
777
778
779
        .route("/v1/conversations", post(v1_conversations_create))
        .route(
            "/v1/conversations/{conversation_id}",
            get(v1_conversations_get)
                .post(v1_conversations_update)
                .delete(v1_conversations_delete),
        )
780
781
        .route(
            "/v1/conversations/{conversation_id}/items",
782
783
784
785
786
            get(v1_conversations_list_items).post(v1_conversations_create_items),
        )
        .route(
            "/v1/conversations/{conversation_id}/items/{item_id}",
            get(v1_conversations_get_item).delete(v1_conversations_delete_item),
787
        )
788
789
        .route_layer(axum::middleware::from_fn_with_state(
            app_state.clone(),
790
            middleware::concurrency_limit_middleware,
791
792
793
794
        ))
        .route_layer(axum::middleware::from_fn_with_state(
            auth_config.clone(),
            middleware::auth_middleware,
795
        ));
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810

    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))
811
812
813
814
815
        .route("/get_loads", get(get_loads))
        .route_layer(axum::middleware::from_fn_with_state(
            auth_config.clone(),
            middleware::auth_middleware,
        ));
816

817
818
819
820
    let worker_routes = Router::new()
        .route("/workers", post(create_worker))
        .route("/workers", get(list_workers_rest))
        .route("/workers/{url}", get(get_worker))
821
822
823
824
825
        .route("/workers/{url}", delete(delete_worker))
        .route_layer(axum::middleware::from_fn_with_state(
            auth_config.clone(),
            middleware::auth_middleware,
        ));
826

827
828
829
830
    Router::new()
        .merge(protected_routes)
        .merge(public_routes)
        .merge(admin_routes)
831
        .merge(worker_routes)
832
        .layer(axum::extract::DefaultBodyLimit::max(max_payload_size))
833
834
835
        .layer(tower_http::limit::RequestBodyLimitLayer::new(
            max_payload_size,
        ))
836
837
        .layer(middleware::create_logging_layer())
        .layer(middleware::RequestIdLayer::new(request_id_headers))
838
839
840
841
842
843
        .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>> {
844
845
846
847
    static LOGGING_INITIALIZED: AtomicBool = AtomicBool::new(false);

    let _log_guard = if !LOGGING_INITIALIZED.swap(true, Ordering::SeqCst) {
        Some(logging::init_logging(LoggingConfig {
848
849
850
851
852
853
            level: config
                .log_level
                .as_deref()
                .and_then(|s| match s.to_uppercase().parse::<Level>() {
                    Ok(l) => Some(l),
                    Err(_) => {
854
                        warn!("Invalid log level string: '{s}'. Defaulting to INFO.");
855
856
857
858
                        None
                    }
                })
                .unwrap_or(Level::INFO),
859
860
861
862
863
864
865
866
867
            json_format: false,
            log_dir: config.log_dir.clone(),
            colorize: true,
            log_file_name: "sgl-router".to_string(),
            log_targets: None,
        }))
    } else {
        None
    };
868

869
870
    if let Some(prometheus_config) = &config.prometheus_config {
        metrics::start_prometheus(prometheus_config.clone());
871
872
    }

873
    info!(
874
875
876
877
878
        "Starting router on {}:{} | mode: {:?} | policy: {:?} | max_payload: {}MB",
        config.host,
        config.port,
        config.router_config.mode,
        config.router_config.policy,
879
880
881
        config.max_payload_size / (1024 * 1024)
    );

882
    let client = Client::builder()
883
        .pool_idle_timeout(Some(Duration::from_secs(50)))
884
        .pool_max_idle_per_host(500)
885
        .timeout(Duration::from_secs(config.request_timeout_secs))
886
        .connect_timeout(Duration::from_secs(10))
887
        .tcp_nodelay(true)
888
        .tcp_keepalive(Some(Duration::from_secs(30)))
889
890
891
        .build()
        .expect("Failed to create HTTP client");

892
    let app_context = AppContext::new(
893
894
895
        config.router_config.clone(),
        client.clone(),
        config.router_config.max_concurrent_requests,
896
        config.router_config.rate_limit_tokens_per_second,
897
898
899
900
    )?;

    let app_context = Arc::new(app_context);

901
902
903
904
    info!(
        "Initializing workers for routing mode: {:?}",
        config.router_config.mode
    );
905
    WorkerManager::initialize_workers(
906
907
908
909
910
911
        &config.router_config,
        &app_context.worker_registry,
        Some(&app_context.policy_registry),
    )
    .await
    .map_err(|e| format!("Failed to initialize workers: {}", e))?;
912
913
914
915
916
917
918

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

919
920
    let router_manager = RouterManager::from_config(&config, &app_context).await?;
    let router: Arc<dyn RouterTrait> = router_manager.clone();
921
922
923
924
925
926
927
928

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

930
931
932
933
934
    if let Some(ref load_monitor) = app_context.load_monitor {
        load_monitor.start().await;
        info!("Started LoadMonitor for PowerOfTwo policies");
    }

935
    let (limiter, processor) = middleware::ConcurrencyLimiter::new(
936
937
938
939
940
        app_context.rate_limiter.clone(),
        config.router_config.queue_size,
        Duration::from_secs(config.router_config.queue_timeout_secs),
    );

941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
    if app_context.rate_limiter.is_none() {
        info!("Rate limiting is disabled (max_concurrent_requests = -1)");
    }

    match processor {
        Some(proc) => {
            spawn(proc.run());
            info!(
                "Started request queue (size: {}, timeout: {}s)",
                config.router_config.queue_size, config.router_config.queue_timeout_secs
            );
        }
        None => {
            info!(
                "Rate limiting enabled (max_concurrent_requests = {}, queue disabled)",
                config.router_config.max_concurrent_requests
            );
        }
959
960
    }

961
    let app_state = Arc::new(AppState {
962
        router,
963
        context: app_context.clone(),
964
        concurrency_queue_tx: limiter.queue_tx.clone(),
965
        router_manager: Some(router_manager),
966
    });
967
968
    if let Some(service_discovery_config) = config.service_discovery_config {
        if service_discovery_config.enabled {
969
970
            let app_context_arc = Arc::clone(&app_state.context);
            match start_service_discovery(service_discovery_config, app_context_arc).await {
971
                Ok(handle) => {
972
                    info!("Service discovery started");
973
974
975
976
977
978
979
                    spawn(async move {
                        if let Err(e) = handle.await {
                            error!("Service discovery task failed: {:?}", e);
                        }
                    });
                }
                Err(e) => {
980
                    error!("Failed to start service discovery: {e}");
981
982
983
984
985
986
                    warn!("Continuing without service discovery");
                }
            }
        }
    }

987
    info!(
988
        "Router ready | workers: {:?}",
989
        WorkerManager::get_worker_urls(&app_state.context.worker_registry)
990
    );
991

992
993
994
995
996
997
998
999
1000
    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(),
        ]
    });

1001
1002
1003
1004
    let auth_config = AuthConfig {
        api_key: config.router_config.api_key.clone(),
    };

1005
1006
    let app = build_app(
        app_state,
1007
        auth_config,
1008
1009
1010
1011
        config.max_payload_size,
        request_id_headers,
        config.router_config.cors_allowed_origins.clone(),
    );
1012

1013
1014
1015
1016
1017
1018
    // TcpListener::bind accepts &str and handles IPv4/IPv6 via ToSocketAddrs
    let bind_addr = format!("{}:{}", config.host, config.port);
    info!("Starting server on {}", bind_addr);
    let listener = TcpListener::bind(&bind_addr)
        .await
        .map_err(|e| format!("Failed to bind to {}: {}", bind_addr, e))?;
1019
    serve(listener, app)
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
        .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))
1078
}