dynamo_runtime/
system_status_server.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4// TODO: (DEP-635) this file should be renamed to system_http_server.rs
5//  it is being used not just for status, health, but others like loras management.
6
7use crate::config::HealthStatus;
8use crate::config::environment_names::logging as env_logging;
9use crate::config::environment_names::runtime::canary as env_canary;
10use crate::config::environment_names::runtime::system as env_system;
11use crate::logging::make_system_request_span;
12use crate::metrics::MetricsHierarchy;
13use crate::traits::DistributedRuntimeProvider;
14use axum::{
15    Router,
16    body::Bytes,
17    extract::{Json, Path, State},
18    http::StatusCode,
19    response::IntoResponse,
20    routing::{any, delete, get, post},
21};
22use futures::StreamExt;
23use serde::{Deserialize, Serialize};
24use serde_json::json;
25use std::collections::HashMap;
26use std::sync::{Arc, OnceLock};
27use std::time::Instant;
28use tokio::{net::TcpListener, task::JoinHandle};
29use tokio_util::sync::CancellationToken;
30use tower_http::trace::TraceLayer;
31
32/// System status server information containing socket address and handle
33#[derive(Debug)]
34pub struct SystemStatusServerInfo {
35    pub socket_addr: std::net::SocketAddr,
36    pub handle: Option<Arc<JoinHandle<()>>>,
37}
38
39impl SystemStatusServerInfo {
40    pub fn new(socket_addr: std::net::SocketAddr, handle: Option<JoinHandle<()>>) -> Self {
41        Self {
42            socket_addr,
43            handle: handle.map(Arc::new),
44        }
45    }
46
47    pub fn address(&self) -> String {
48        self.socket_addr.to_string()
49    }
50
51    pub fn hostname(&self) -> String {
52        self.socket_addr.ip().to_string()
53    }
54
55    pub fn port(&self) -> u16 {
56        self.socket_addr.port()
57    }
58}
59
60impl Clone for SystemStatusServerInfo {
61    fn clone(&self) -> Self {
62        Self {
63            socket_addr: self.socket_addr,
64            handle: self.handle.clone(),
65        }
66    }
67}
68
69/// System status server state containing the distributed runtime reference
70pub struct SystemStatusState {
71    // global drt registry is for printing out the entire Prometheus format output
72    root_drt: Arc<crate::DistributedRuntime>,
73    // Discovery metadata (only for Kubernetes backend)
74    discovery_metadata: Option<Arc<tokio::sync::RwLock<crate::discovery::DiscoveryMetadata>>>,
75}
76
77impl SystemStatusState {
78    /// Create new system status server state with the provided distributed runtime
79    pub fn new(
80        drt: Arc<crate::DistributedRuntime>,
81        discovery_metadata: Option<Arc<tokio::sync::RwLock<crate::discovery::DiscoveryMetadata>>>,
82    ) -> anyhow::Result<Self> {
83        Ok(Self {
84            root_drt: drt,
85            discovery_metadata,
86        })
87    }
88
89    /// Get a reference to the distributed runtime
90    pub fn drt(&self) -> &crate::DistributedRuntime {
91        &self.root_drt
92    }
93
94    /// Get a reference to the discovery metadata if available
95    pub fn discovery_metadata(
96        &self,
97    ) -> Option<&Arc<tokio::sync::RwLock<crate::discovery::DiscoveryMetadata>>> {
98        self.discovery_metadata.as_ref()
99    }
100}
101
102/// Request body for POST /v1/loras
103#[derive(Debug, Clone, Deserialize, Serialize)]
104pub struct LoadLoraRequest {
105    pub lora_name: String,
106    pub source: LoraSource,
107}
108
109/// Source information for loading a LoRA
110#[derive(Debug, Clone, Deserialize, Serialize)]
111pub struct LoraSource {
112    pub uri: String,
113}
114
115/// Response body for LoRA operations
116#[derive(Debug, Clone, Deserialize, Serialize)]
117pub struct LoraResponse {
118    pub status: String,
119    #[serde(skip_serializing_if = "Option::is_none")]
120    pub message: Option<String>,
121    #[serde(skip_serializing_if = "Option::is_none")]
122    pub lora_name: Option<String>,
123    #[serde(skip_serializing_if = "Option::is_none")]
124    pub lora_id: Option<u64>,
125    #[serde(skip_serializing_if = "Option::is_none")]
126    pub loras: Option<serde_json::Value>,
127    #[serde(skip_serializing_if = "Option::is_none")]
128    pub count: Option<usize>,
129}
130
131/// Start system status server with metrics support
132pub async fn spawn_system_status_server(
133    host: &str,
134    port: u16,
135    cancel_token: CancellationToken,
136    drt: Arc<crate::DistributedRuntime>,
137    discovery_metadata: Option<Arc<tokio::sync::RwLock<crate::discovery::DiscoveryMetadata>>>,
138) -> anyhow::Result<(std::net::SocketAddr, tokio::task::JoinHandle<()>)> {
139    // Create system status server state with the provided distributed runtime
140    let server_state = Arc::new(SystemStatusState::new(drt, discovery_metadata)?);
141    let health_path = server_state
142        .drt()
143        .system_health()
144        .lock()
145        .health_path()
146        .to_string();
147    let live_path = server_state
148        .drt()
149        .system_health()
150        .lock()
151        .live_path()
152        .to_string();
153
154    // Check if LoRA feature is enabled
155    let lora_enabled =
156        crate::config::env_is_truthy(crate::config::environment_names::llm::DYN_LORA_ENABLED);
157
158    let mut app = Router::new()
159        .route(
160            &health_path,
161            get({
162                let state = Arc::clone(&server_state);
163                move || health_handler(state)
164            }),
165        )
166        .route(
167            &live_path,
168            get({
169                let state = Arc::clone(&server_state);
170                move || health_handler(state)
171            }),
172        )
173        .route(
174            "/metrics",
175            get({
176                let state = Arc::clone(&server_state);
177                move || metrics_handler(state)
178            }),
179        )
180        .route(
181            "/metadata",
182            get({
183                let state = Arc::clone(&server_state);
184                move || metadata_handler(state)
185            }),
186        )
187        .route(
188            "/engine/{*path}",
189            any({
190                let state = Arc::clone(&server_state);
191                move |path, body| engine_route_handler(state, path, body)
192            }),
193        );
194
195    // Add LoRA routes only if DYN_LORA_ENABLED is set to true
196    if lora_enabled {
197        app = app
198            .route(
199                "/v1/loras",
200                get({
201                    let state = Arc::clone(&server_state);
202                    move || list_loras_handler(State(state))
203                })
204                .post({
205                    let state = Arc::clone(&server_state);
206                    move |body| load_lora_handler(State(state), body)
207                }),
208            )
209            .route(
210                "/v1/loras/{*lora_name}",
211                delete({
212                    let state = Arc::clone(&server_state);
213                    move |path| unload_lora_handler(State(state), path)
214                }),
215            );
216    }
217
218    // Self-hosted MDC files. Always mounted; empty registry → 404.
219    // The endpoint triple disambiguates multi-LocalModel-per-DRT; the
220    // suffix segment (LoRA slug or `_base`) scopes per-registration so
221    // detaching one doesn't wipe another's entries.
222    app = app.route(
223        "/v1/metadata/{namespace}/{component}/{endpoint}/{model_slug}/{model_suffix}/{*filename}",
224        get({
225            let state = Arc::clone(&server_state);
226            move |path| metadata_file_handler(State(state), path)
227        }),
228    );
229
230    let app = app
231        .fallback(|| async {
232            tracing::info!("[fallback handler] called");
233            (StatusCode::NOT_FOUND, "Route not found").into_response()
234        })
235        .layer(TraceLayer::new_for_http().make_span_with(make_system_request_span));
236
237    let address = format!("{}:{}", host, port);
238    tracing::info!("[spawn_system_status_server] binding to: {address}");
239
240    let listener = match TcpListener::bind(&address).await {
241        Ok(listener) => {
242            // get the actual address and port, print in debug level
243            let actual_address = listener.local_addr()?;
244            tracing::info!(
245                "[spawn_system_status_server] system status server bound to: {}",
246                actual_address
247            );
248            (listener, actual_address)
249        }
250        Err(e) => {
251            tracing::error!("Failed to bind to address {}: {}", address, e);
252            return Err(anyhow::anyhow!("Failed to bind to address: {}", e));
253        }
254    };
255    let (listener, actual_address) = listener;
256
257    let observer = cancel_token.child_token();
258    // Spawn the server in the background and return the handle
259    let handle = tokio::spawn(async move {
260        if let Err(e) = axum::serve(listener, app)
261            .with_graceful_shutdown(observer.cancelled_owned())
262            .await
263        {
264            tracing::error!("System status server error: {e}");
265        }
266    });
267
268    Ok((actual_address, handle))
269}
270
271/// Health handler with optional active health checking
272#[tracing::instrument(skip_all, level = "trace")]
273async fn health_handler(state: Arc<SystemStatusState>) -> impl IntoResponse {
274    // Get basic health status
275    let system_health = state.drt().system_health();
276    let system_health_lock = system_health.lock();
277    let (healthy, endpoints) = system_health_lock.get_health_status();
278    let uptime = Some(system_health_lock.uptime());
279    drop(system_health_lock);
280
281    let healthy_string = if healthy { "ready" } else { "notready" };
282    let status_code = if healthy {
283        StatusCode::OK
284    } else {
285        StatusCode::SERVICE_UNAVAILABLE
286    };
287
288    let response = json!({
289        "status": healthy_string,
290        "uptime": uptime,
291        "endpoints": endpoints,
292    });
293
294    tracing::trace!("Response {}", response.to_string());
295
296    (status_code, response.to_string())
297}
298
299/// Metrics handler with DistributedRuntime uptime
300#[tracing::instrument(skip_all, level = "trace")]
301async fn metrics_handler(state: Arc<SystemStatusState>) -> impl IntoResponse {
302    // Get all metrics from the DistributedRuntime.
303    // The uptime gauge is updated automatically via a PrometheusUpdateCallback
304    // registered in DistributedRuntime::new(), so it is always fresh before scrape.
305    //
306    // NOTE: We use a multi-registry model (e.g. one registry per endpoint) and merge at scrape time,
307    // so /metrics traverses registered child registries and produces a single combined output.
308    let response = match state.drt().metrics().prometheus_expfmt() {
309        Ok(r) => r,
310        Err(e) => {
311            tracing::error!("Failed to get metrics from registry: {e}");
312            return (
313                StatusCode::INTERNAL_SERVER_ERROR,
314                "Failed to get metrics".to_string(),
315            );
316        }
317    };
318
319    (StatusCode::OK, response)
320}
321
322/// Metadata handler
323#[tracing::instrument(skip_all, level = "trace")]
324async fn metadata_handler(state: Arc<SystemStatusState>) -> impl IntoResponse {
325    // Check if discovery metadata is available
326    let metadata = match state.discovery_metadata() {
327        Some(metadata) => metadata,
328        None => {
329            tracing::debug!("Metadata endpoint called but no discovery metadata available");
330            return (
331                StatusCode::NOT_FOUND,
332                "Discovery metadata not available".to_string(),
333            )
334                .into_response();
335        }
336    };
337
338    // Read the metadata
339    let metadata_guard = metadata.read().await;
340
341    // Serialize to JSON
342    match serde_json::to_string(&*metadata_guard) {
343        Ok(json) => {
344            tracing::trace!("Returning metadata: {} bytes", json.len());
345            (StatusCode::OK, json).into_response()
346        }
347        Err(e) => {
348            tracing::error!("Failed to serialize metadata: {e}");
349            (
350                StatusCode::INTERNAL_SERVER_ERROR,
351                "Failed to serialize metadata".to_string(),
352            )
353                .into_response()
354        }
355    }
356}
357
358/// Handler for POST /v1/loras - Load a LoRA adapter
359#[tracing::instrument(skip_all, level = "debug")]
360async fn load_lora_handler(
361    State(state): State<Arc<SystemStatusState>>,
362    Json(request): Json<LoadLoraRequest>,
363) -> impl IntoResponse {
364    tracing::info!("Loading LoRA: {}", request.lora_name);
365
366    // Call the load_lora endpoint for each available backend
367    match call_lora_endpoint(
368        state.drt(),
369        "load_lora",
370        json!({
371            "lora_name": request.lora_name,
372            "source": {
373                "uri": request.source.uri
374            },
375        }),
376    )
377    .await
378    {
379        Ok(response) => {
380            if response.status == "error" {
381                tracing::error!(
382                    "Failed to load LoRA {}: {}",
383                    request.lora_name,
384                    response.message.as_deref().unwrap_or("Unknown error")
385                );
386                (StatusCode::INTERNAL_SERVER_ERROR, Json(response))
387            } else {
388                tracing::info!("LoRA loaded successfully: {}", request.lora_name);
389                (StatusCode::OK, Json(response))
390            }
391        }
392        Err(e) => {
393            tracing::error!("Failed to load LoRA {}: {}", request.lora_name, e);
394            (
395                StatusCode::INTERNAL_SERVER_ERROR,
396                Json(LoraResponse {
397                    status: "error".to_string(),
398                    message: Some(e.to_string()),
399                    lora_name: Some(request.lora_name),
400                    lora_id: None,
401                    loras: None,
402                    count: None,
403                }),
404            )
405        }
406    }
407}
408
409/// Handler for DELETE /v1/loras/*lora_name - Unload a LoRA adapter
410#[tracing::instrument(skip_all, level = "debug")]
411async fn unload_lora_handler(
412    State(state): State<Arc<SystemStatusState>>,
413    Path(lora_name): Path<String>,
414) -> impl IntoResponse {
415    // Strip the leading slash from the wildcard capture
416    let lora_name = lora_name
417        .strip_prefix('/')
418        .unwrap_or(&lora_name)
419        .to_string();
420    tracing::info!("Unloading LoRA: {lora_name}");
421
422    // Call the unload_lora endpoint for each available backend
423    match call_lora_endpoint(
424        state.drt(),
425        "unload_lora",
426        json!({
427            "lora_name": lora_name.clone(),
428        }),
429    )
430    .await
431    {
432        Ok(response) => {
433            if response.status == "error" {
434                tracing::error!(
435                    "Failed to unload LoRA {}: {}",
436                    lora_name,
437                    response.message.as_deref().unwrap_or("Unknown error")
438                );
439                (StatusCode::INTERNAL_SERVER_ERROR, Json(response))
440            } else {
441                tracing::info!("LoRA unloaded successfully: {lora_name}");
442                (StatusCode::OK, Json(response))
443            }
444        }
445        Err(e) => {
446            tracing::error!("Failed to unload LoRA {}: {}", lora_name, e);
447            (
448                StatusCode::INTERNAL_SERVER_ERROR,
449                Json(LoraResponse {
450                    status: "error".to_string(),
451                    message: Some(e.to_string()),
452                    lora_name: Some(lora_name),
453                    lora_id: None,
454                    loras: None,
455                    count: None,
456                }),
457            )
458        }
459    }
460}
461
462/// Handler for GET /v1/loras - List all LoRA adapters
463#[tracing::instrument(skip_all, level = "debug")]
464async fn list_loras_handler(State(state): State<Arc<SystemStatusState>>) -> impl IntoResponse {
465    tracing::info!("Listing all LoRAs");
466
467    // Call the list_loras endpoint for each available backend
468    match call_lora_endpoint(state.drt(), "list_loras", json!({})).await {
469        Ok(response) => {
470            tracing::info!("Successfully retrieved LoRA list");
471            (StatusCode::OK, Json(response))
472        }
473        Err(e) => {
474            tracing::error!("Failed to list LoRAs: {e}");
475            (
476                StatusCode::INTERNAL_SERVER_ERROR,
477                Json(LoraResponse {
478                    status: "error".to_string(),
479                    message: Some(e.to_string()),
480                    lora_name: None,
481                    lora_id: None,
482                    loras: None,
483                    count: None,
484                }),
485            )
486        }
487    }
488}
489
490/// `GET /v1/metadata/{namespace}/{component}/{endpoint}/{slug}/{suffix}/{filename}`
491/// — 404 on miss, 500 on read error, raw bytes on hit. Consumer blake3-verifies.
492async fn metadata_file_handler(
493    State(state): State<Arc<SystemStatusState>>,
494    Path((namespace, component, endpoint, model_slug, model_suffix, filename)): Path<(
495        String,
496        String,
497        String,
498        String,
499        String,
500        String,
501    )>,
502) -> impl IntoResponse {
503    let path = match state.drt().metadata_artifacts().get(
504        &namespace,
505        &component,
506        &endpoint,
507        &model_slug,
508        &model_suffix,
509        &filename,
510    ) {
511        Some(p) => p,
512        None => {
513            tracing::debug!(
514                namespace,
515                component,
516                endpoint,
517                model_slug,
518                model_suffix,
519                filename,
520                "metadata artifact not registered for self-host"
521            );
522            return (StatusCode::NOT_FOUND, "Not found").into_response();
523        }
524    };
525
526    match tokio::fs::read(&path).await {
527        Ok(bytes) => (StatusCode::OK, bytes).into_response(),
528        Err(err) => {
529            tracing::error!(
530                namespace,
531                component,
532                endpoint,
533                model_slug,
534                model_suffix,
535                filename,
536                path = %path.display(),
537                %err,
538                "failed to read self-hosted metadata file"
539            );
540            (StatusCode::INTERNAL_SERVER_ERROR, "Failed to read file").into_response()
541        }
542    }
543}
544
545/// Helper function to call a LoRA management endpoint for the local worker.
546///
547/// Resolution order (both are in-process, never network discovery):
548/// 1. The legacy local endpoint registry, populated by non-unified workers via
549///    `.register_local_engine()`.
550/// 2. The generic engine-route registry (`/engine/*`). Unified-backend workers
551///    advertise LoRA lifecycle ops (`load_lora`/`unload_lora`/`list_loras`) as
552///    engine *updates*, registered under `update/<name>`; this fallback maps the
553///    bare LoRA name onto that key so the legacy `/v1/loras` surface forwards to
554///    them (the `/v1/loras` compatibility shim).
555///
556/// Because legacy workers populate the local registry, they never reach the
557/// fallback — their `/v1/loras` behavior is unchanged. If neither registry
558/// holds the name, returns an explicit "LoRA management not available" error
559/// rather than an opaque "endpoint not found".
560async fn call_lora_endpoint(
561    drt: &crate::DistributedRuntime,
562    endpoint_name: &str,
563    request_body: serde_json::Value,
564) -> anyhow::Result<LoraResponse> {
565    use crate::engine::AsyncEngine;
566
567    tracing::debug!("Calling LoRA endpoint: '{endpoint_name}'");
568
569    // 1. Legacy local registry (in-process call only).
570    if let Some(engine) = drt.local_endpoint_registry().get(endpoint_name) {
571        tracing::debug!(
572            "Found endpoint '{}' in local registry, calling directly",
573            endpoint_name
574        );
575
576        let request = crate::pipeline::SingleIn::new(request_body);
577        let mut stream = engine.generate(request).await?;
578
579        if let Some(response) = stream.next().await {
580            let response_data = response.data.unwrap_or_default();
581            let lora_response = serde_json::from_value::<LoraResponse>(response_data.clone())
582                .unwrap_or_else(|_| parse_lora_response(&response_data));
583            return Ok(lora_response);
584        }
585
586        anyhow::bail!("No response received from endpoint '{}'", endpoint_name)
587    }
588
589    // 2. Unified-backend engine-update registry fallback. The unified Worker
590    //    registers LoRA ops as engine updates under `update/<name>`, so map the
591    //    bare LoRA endpoint name onto that namespaced key.
592    let update_key = format!("update/{endpoint_name}");
593    if let Some(callback) = drt.engine_routes().get(&update_key) {
594        tracing::debug!(
595            "Found '{}' in engine routes registry, invoking update callback",
596            update_key
597        );
598        let response_data = callback(request_body).await?;
599        let lora_response = serde_json::from_value::<LoraResponse>(response_data.clone())
600            .unwrap_or_else(|_| parse_lora_response(&response_data));
601        return Ok(lora_response);
602    }
603
604    anyhow::bail!(
605        "LoRA management not available: no '{}' handler is registered \
606         (neither a local LoRA endpoint nor an engine update). This worker \
607         either has LoRA disabled or its backend does not support LoRA \
608         management.",
609        endpoint_name
610    )
611}
612
613/// Helper to parse response data into LoraResponse
614fn parse_lora_response(response_data: &serde_json::Value) -> LoraResponse {
615    LoraResponse {
616        status: response_data
617            .get("status")
618            .and_then(|s| s.as_str())
619            .unwrap_or("success")
620            .to_string(),
621        message: response_data
622            .get("message")
623            .and_then(|m| m.as_str())
624            .map(|s| s.to_string()),
625        lora_name: response_data
626            .get("lora_name")
627            .and_then(|n| n.as_str())
628            .map(|s| s.to_string()),
629        lora_id: response_data.get("lora_id").and_then(|id| id.as_u64()),
630        loras: response_data.get("loras").cloned(),
631        count: response_data
632            .get("count")
633            .and_then(|c| c.as_u64())
634            .map(|c| c as usize),
635    }
636}
637
638/// Engine route handler for /engine/* routes
639///
640/// This handler looks up registered callbacks in the engine routes registry
641/// and invokes them with the request body, returning the response as JSON.
642#[tracing::instrument(skip_all, level = "trace", fields(path = %path))]
643async fn engine_route_handler(
644    state: Arc<SystemStatusState>,
645    Path(path): Path<String>,
646    body: Bytes,
647) -> impl IntoResponse {
648    tracing::trace!("Engine route request to /engine/{path}");
649
650    // Parse body as JSON (empty object for GET/empty body)
651    let body_json: serde_json::Value = if body.is_empty() {
652        serde_json::json!({})
653    } else {
654        match serde_json::from_slice(&body) {
655            Ok(json) => json,
656            Err(e) => {
657                tracing::warn!("Invalid JSON in request body: {e}");
658                return (
659                    StatusCode::BAD_REQUEST,
660                    json!({
661                        "error": "Invalid JSON",
662                        "message": format!("{}", e)
663                    })
664                    .to_string(),
665                )
666                    .into_response();
667            }
668        }
669    };
670
671    // Look up callback
672    let callback = match state.drt().engine_routes().get(&path) {
673        Some(cb) => cb,
674        None => {
675            tracing::debug!("Route /engine/{path} not found");
676            return (
677                StatusCode::NOT_FOUND,
678                json!({
679                    "error": "Route not found",
680                    "message": format!("Route /engine/{} not found", path)
681                })
682                .to_string(),
683            )
684                .into_response();
685        }
686    };
687
688    // Call callback (it's async, so await it)
689    match callback(body_json).await {
690        Ok(response) => {
691            tracing::trace!("Engine route handler succeeded for /engine/{path}");
692            (StatusCode::OK, response.to_string()).into_response()
693        }
694        Err(e) => {
695            tracing::error!("Engine route handler error for /engine/{}: {}", path, e);
696            (
697                StatusCode::INTERNAL_SERVER_ERROR,
698                json!({
699                    "error": "Handler error",
700                    "message": format!("{}", e)
701                })
702                .to_string(),
703            )
704                .into_response()
705        }
706    }
707}
708
709// Regular tests: cargo test system_status_server --lib
710#[cfg(test)]
711mod tests {
712    use super::*;
713    use tokio::time::Duration;
714
715    // This is a basic test to verify the HTTP server is working before testing other more complicated tests
716    #[tokio::test]
717    async fn test_http_server_lifecycle() {
718        let cancel_token = CancellationToken::new();
719        let cancel_token_for_server = cancel_token.clone();
720
721        // Test basic HTTP server lifecycle without DistributedRuntime
722        let app = Router::new().route("/test", get(|| async { (StatusCode::OK, "test") }));
723
724        // start HTTP server
725        let server_handle = tokio::spawn(async move {
726            let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
727            let _ = axum::serve(listener, app)
728                .with_graceful_shutdown(cancel_token_for_server.cancelled_owned())
729                .await;
730        });
731
732        // server starts immediately, no need to wait
733
734        // cancel token
735        cancel_token.cancel();
736
737        // wait for the server to shut down
738        let result = tokio::time::timeout(Duration::from_secs(5), server_handle).await;
739        assert!(
740            result.is_ok(),
741            "HTTP server should shut down when cancel token is cancelled"
742        );
743    }
744}
745
746// Integration tests: cargo test system_status_server --lib --features integration
747#[cfg(all(test, feature = "integration"))]
748mod integration_tests {
749    use super::*;
750    use crate::config::environment_names::logging as env_logging;
751    use crate::config::environment_names::runtime::canary as env_canary;
752    use crate::distributed::distributed_test_utils::create_test_drt_async;
753    use crate::metrics::MetricsHierarchy;
754    use anyhow::Result;
755    use rstest::rstest;
756    use std::sync::Arc;
757    use tokio::time::Duration;
758
759    #[tokio::test]
760    async fn test_uptime_from_system_health() {
761        // Test that uptime is available from SystemHealth
762        temp_env::async_with_vars([(env_system::DYN_SYSTEM_PORT, None::<&str>)], async {
763            let drt = create_test_drt_async().await;
764
765            // Get uptime from SystemHealth
766            let uptime = drt.system_health().lock().uptime();
767            // Uptime should exist (even if close to zero)
768            assert!(uptime.as_nanos() > 0 || uptime.is_zero());
769
770            // Sleep briefly and check uptime increases
771            tokio::time::sleep(std::time::Duration::from_millis(100)).await;
772            let uptime_after = drt.system_health().lock().uptime();
773            assert!(uptime_after > uptime);
774        })
775        .await;
776    }
777
778    #[tokio::test]
779    async fn test_runtime_metrics_initialization_and_namespace() {
780        // Test that metrics have correct namespace
781        temp_env::async_with_vars([(env_system::DYN_SYSTEM_PORT, None::<&str>)], async {
782            let drt = create_test_drt_async().await;
783            // SystemStatusState is already created in distributed.rs
784            // so we don't need to create it again here
785
786            // The uptime_seconds metric should already be registered and available
787            let response = drt.metrics().prometheus_expfmt().unwrap();
788            println!("Full metrics response:\n{}", response);
789
790            // Check that uptime_seconds metric is present with correct namespace
791            assert!(
792                response.contains("# HELP dynamo_component_uptime_seconds"),
793                "Should contain uptime_seconds help text"
794            );
795            assert!(
796                response.contains("# TYPE dynamo_component_uptime_seconds gauge"),
797                "Should contain uptime_seconds type"
798            );
799            assert!(
800                response.contains("dynamo_component_uptime_seconds"),
801                "Should contain uptime_seconds metric with correct namespace"
802            );
803        })
804        .await;
805    }
806
807    #[tokio::test]
808    async fn test_uptime_gauge_updates() {
809        // Test that the uptime gauge is properly updated and increases over time
810        temp_env::async_with_vars([(env_system::DYN_SYSTEM_PORT, None::<&str>)], async {
811            let drt = create_test_drt_async().await;
812
813            // Get initial uptime
814            let initial_uptime = drt.system_health().lock().uptime();
815
816            // Update the gauge with initial value
817            drt.system_health().lock().update_uptime_gauge();
818
819            // Sleep for 100ms
820            tokio::time::sleep(std::time::Duration::from_millis(100)).await;
821
822            // Get uptime after sleep
823            let uptime_after_sleep = drt.system_health().lock().uptime();
824
825            // Update the gauge again
826            drt.system_health().lock().update_uptime_gauge();
827
828            // Verify uptime increased by at least 100ms
829            let elapsed = uptime_after_sleep - initial_uptime;
830            assert!(
831                elapsed >= std::time::Duration::from_millis(100),
832                "Uptime should have increased by at least 100ms after sleep, but only increased by {:?}",
833                elapsed
834            );
835        })
836        .await;
837    }
838
839    #[tokio::test]
840    async fn test_http_requests_fail_when_system_disabled() {
841        // Test that system status server is not running when disabled
842        temp_env::async_with_vars([(env_system::DYN_SYSTEM_PORT, None::<&str>)], async {
843            let drt = create_test_drt_async().await;
844
845            // Verify that system status server info is None when disabled
846            let system_info = drt.system_status_server_info();
847            assert!(
848                system_info.is_none(),
849                "System status server should not be running when disabled"
850            );
851
852            println!("✓ System status server correctly disabled when not enabled");
853        })
854        .await;
855    }
856
857    /// This test verifies the health and liveness endpoints of the system status server.
858    /// It checks that the endpoints respond with the correct HTTP status codes and bodies
859    /// based on the initial health status and any custom endpoint paths provided via environment variables.
860    /// The test is parameterized using multiple #[case] attributes to cover various scenarios,
861    /// including different initial health states ("ready" and "notready"), default and custom endpoint paths,
862    /// and expected response codes and bodies.
863    #[rstest]
864    #[case("ready", 200, "ready", None, None, 3)]
865    #[case("notready", 503, "notready", None, None, 3)]
866    #[case("ready", 200, "ready", Some("/custom/health"), Some("/custom/live"), 5)]
867    #[case(
868        "notready",
869        503,
870        "notready",
871        Some("/custom/health"),
872        Some("/custom/live"),
873        5
874    )]
875    #[tokio::test]
876    #[cfg(feature = "integration")]
877    async fn test_health_endpoints(
878        #[case] starting_health_status: &'static str,
879        #[case] expected_status: u16,
880        #[case] expected_body: &'static str,
881        #[case] custom_health_path: Option<&'static str>,
882        #[case] custom_live_path: Option<&'static str>,
883        #[case] expected_num_tests: usize,
884    ) {
885        use std::sync::Arc;
886        // use tokio::io::{AsyncReadExt, AsyncWriteExt};
887        // use reqwest for HTTP requests
888
889        // Closure call is needed here to satisfy async_with_vars
890
891        crate::logging::init();
892
893        #[allow(clippy::redundant_closure_call)]
894        temp_env::async_with_vars(
895            [
896                (env_system::DYN_SYSTEM_PORT, Some("0")),
897                (
898                    env_system::DYN_SYSTEM_STARTING_HEALTH_STATUS,
899                    Some(starting_health_status),
900                ),
901                (env_system::DYN_SYSTEM_HEALTH_PATH, custom_health_path),
902                (env_system::DYN_SYSTEM_LIVE_PATH, custom_live_path),
903            ],
904            (async || {
905                let drt = Arc::new(create_test_drt_async().await);
906
907                // Get system status server info from DRT (instead of manually spawning)
908                let system_info = drt
909                    .system_status_server_info()
910                    .expect("System status server should be started by DRT");
911                let addr = system_info.socket_addr;
912
913                let client = reqwest::Client::new();
914
915                // Prepare test cases
916                let mut test_cases = vec![];
917                match custom_health_path {
918                    None => {
919                        // When using default paths, test the default paths
920                        test_cases.push(("/health", expected_status, expected_body));
921                    }
922                    Some(chp) => {
923                        // When using custom paths, default paths should not exist
924                        test_cases.push(("/health", 404, "Route not found"));
925                        test_cases.push((chp, expected_status, expected_body));
926                    }
927                }
928                match custom_live_path {
929                    None => {
930                        // When using default paths, test the default paths
931                        test_cases.push(("/live", expected_status, expected_body));
932                    }
933                    Some(clp) => {
934                        // When using custom paths, default paths should not exist
935                        test_cases.push(("/live", 404, "Route not found"));
936                        test_cases.push((clp, expected_status, expected_body));
937                    }
938                }
939                test_cases.push(("/someRandomPathNotFoundHere", 404, "Route not found"));
940                assert_eq!(test_cases.len(), expected_num_tests);
941
942                for (path, expect_status, expect_body) in test_cases {
943                    println!("[test] Sending request to {}", path);
944                    let url = format!("http://{}{}", addr, path);
945                    let response = client.get(&url).send().await.unwrap();
946                    let status = response.status();
947                    let body = response.text().await.unwrap();
948                    println!(
949                        "[test] Response for {}: status={}, body={:?}",
950                        path, status, body
951                    );
952                    assert_eq!(
953                        status, expect_status,
954                        "Response: status={}, body={:?}",
955                        status, body
956                    );
957                    assert!(
958                        body.contains(expect_body),
959                        "Response: status={}, body={:?}",
960                        status,
961                        body
962                    );
963                }
964            })(),
965        )
966        .await;
967    }
968
969    #[tokio::test]
970    async fn test_health_endpoint_tracing() -> Result<()> {
971        use std::sync::Arc;
972
973        // Closure call is needed here to satisfy async_with_vars
974
975        #[allow(clippy::redundant_closure_call)]
976        let _ = temp_env::async_with_vars(
977            [
978                (env_system::DYN_SYSTEM_PORT, Some("0")),
979                (env_system::DYN_SYSTEM_STARTING_HEALTH_STATUS, Some("ready")),
980                (env_logging::DYN_LOGGING_JSONL, Some("1")),
981                (env_logging::DYN_LOG, Some("trace")),
982            ],
983            (async || {
984                // TODO Add proper testing for
985                // trace id and parent id
986
987                crate::logging::init();
988
989                let drt = Arc::new(create_test_drt_async().await);
990
991                // Get system status server info from DRT (instead of manually spawning)
992                let system_info = drt
993                    .system_status_server_info()
994                    .expect("System status server should be started by DRT");
995                let addr = system_info.socket_addr;
996                let client = reqwest::Client::new();
997                for path in [("/health"), ("/live"), ("/someRandomPathNotFoundHere")] {
998                    let traceparent_value =
999                        "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01";
1000                    let tracestate_value = "vendor1=opaqueValue1,vendor2=opaqueValue2";
1001                    let mut headers = reqwest::header::HeaderMap::new();
1002                    headers.insert(
1003                        reqwest::header::HeaderName::from_static("traceparent"),
1004                        reqwest::header::HeaderValue::from_str(traceparent_value)?,
1005                    );
1006                    headers.insert(
1007                        reqwest::header::HeaderName::from_static("tracestate"),
1008                        reqwest::header::HeaderValue::from_str(tracestate_value)?,
1009                    );
1010                    let url = format!("http://{}{}", addr, path);
1011                    let response = client.get(&url).headers(headers).send().await.unwrap();
1012                    let status = response.status();
1013                    let body = response.text().await.unwrap();
1014                    tracing::info!(body = body, status = status.to_string());
1015                }
1016
1017                Ok::<(), anyhow::Error>(())
1018            })(),
1019        )
1020        .await;
1021        Ok(())
1022    }
1023
1024    #[tokio::test]
1025    async fn test_health_endpoint_with_changing_health_status() {
1026        // Test health endpoint starts in not ready status, then becomes ready
1027        // when endpoints are created (DYN_SYSTEM_USE_ENDPOINT_HEALTH_STATUS=generate)
1028        const ENDPOINT_NAME: &str = "generate";
1029        const ENDPOINT_HEALTH_CONFIG: &str = "[\"generate\"]";
1030        temp_env::async_with_vars(
1031            [
1032                (env_system::DYN_SYSTEM_PORT, Some("0")),
1033                (env_system::DYN_SYSTEM_STARTING_HEALTH_STATUS, Some("notready")),
1034                (env_system::DYN_SYSTEM_USE_ENDPOINT_HEALTH_STATUS, Some(ENDPOINT_HEALTH_CONFIG)),
1035            ],
1036            async {
1037                let drt = Arc::new(create_test_drt_async().await);
1038
1039                // Check if system status server was started
1040                let system_info_opt = drt.system_status_server_info();
1041
1042                // Ensure system status server was spawned by DRT
1043                assert!(
1044                    system_info_opt.is_some(),
1045                    "System status server was not spawned by DRT. Expected DRT to spawn server when DYN_SYSTEM_PORT is set to a positive value, but system_status_server_info() returned None. Environment: DYN_SYSTEM_PORT={:?}",
1046                    std::env::var(env_system::DYN_SYSTEM_PORT)
1047                );
1048
1049                // Get the system status server info from DRT - this should never fail now due to above check
1050                let system_info = system_info_opt.unwrap();
1051                let addr = system_info.socket_addr;
1052
1053                // Initially check health - should be not ready
1054                let client = reqwest::Client::new();
1055                let health_url = format!("http://{}/health", addr);
1056
1057                let response = client.get(&health_url).send().await.unwrap();
1058                let status = response.status();
1059                let body = response.text().await.unwrap();
1060
1061                // Health should be not ready (503) initially
1062                assert_eq!(status, 503, "Health should be 503 (not ready) initially, got: {}", status);
1063                assert!(body.contains("\"status\":\"notready\""), "Health should contain status notready");
1064
1065                // Now create a namespace, component, and endpoint to make the system healthy
1066                let namespace = drt.namespace("ns1234").unwrap();
1067                let component = namespace.component("comp1234").unwrap();
1068
1069                // Create a simple test handler
1070                use crate::pipeline::{async_trait, network::Ingress, AsyncEngine, AsyncEngineContextProvider, Error, ManyOut, SingleIn};
1071                use crate::protocols::annotated::Annotated;
1072
1073                struct TestHandler;
1074
1075                #[async_trait]
1076                impl AsyncEngine<SingleIn<String>, ManyOut<Annotated<String>>, anyhow::Error> for TestHandler {
1077                    async fn generate(&self, input: SingleIn<String>) -> anyhow::Result<ManyOut<Annotated<String>>> {
1078                        let (data, ctx) = input.into_parts();
1079                        let response = Annotated::from_data(format!("You responded: {}", data));
1080                        Ok(crate::pipeline::ResponseStream::new(
1081                            Box::pin(crate::stream::iter(vec![response])),
1082                            ctx.context()
1083                        ))
1084                    }
1085                }
1086
1087                // Create the ingress and start the endpoint service
1088                let ingress = Ingress::for_engine(std::sync::Arc::new(TestHandler)).unwrap();
1089
1090                // Start the service and endpoint with a health check payload
1091                // This will automatically register the endpoint for health monitoring
1092                tokio::spawn(async move {
1093                    let _ = component.endpoint(ENDPOINT_NAME)
1094                        .endpoint_builder()
1095                        .handler(ingress)
1096                        .health_check_payload(serde_json::json!({
1097                            "test": "health_check"
1098                        }))
1099                        .start()
1100                        .await;
1101                });
1102
1103                // Hit health endpoint 200 times to verify consistency
1104                let mut success_count = 0;
1105                let mut failures = Vec::new();
1106
1107                for i in 1..=200 {
1108                    let response = client.get(&health_url).send().await.unwrap();
1109                    let status = response.status();
1110                    let body = response.text().await.unwrap();
1111
1112                    if status == 200 && body.contains("\"status\":\"ready\"") {
1113                        success_count += 1;
1114                    } else {
1115                        failures.push((i, status.as_u16(), body.clone()));
1116                        if failures.len() <= 5 {  // Only log first 5 failures
1117                            tracing::warn!("Request {}: status={}, body={}", i, status, body);
1118                        }
1119                    }
1120                }
1121
1122                tracing::info!("Health endpoint test results: {success_count}/200 requests succeeded");
1123                if !failures.is_empty() {
1124                    tracing::warn!("Failed requests: {}", failures.len());
1125                }
1126
1127                // Expect at least 150 out of 200 requests to be successful
1128                assert!(success_count >= 150, "Expected at least 150 out of 200 requests to succeed, but only {} succeeded", success_count);
1129            },
1130        )
1131        .await;
1132    }
1133
1134    #[tokio::test]
1135    async fn test_spawn_system_status_server_endpoints() {
1136        // use reqwest for HTTP requests
1137        temp_env::async_with_vars(
1138            [
1139                (env_system::DYN_SYSTEM_PORT, Some("0")),
1140                (env_system::DYN_SYSTEM_STARTING_HEALTH_STATUS, Some("ready")),
1141            ],
1142            async {
1143                let drt = Arc::new(create_test_drt_async().await);
1144
1145                // Get system status server info from DRT (instead of manually spawning)
1146                let system_info = drt
1147                    .system_status_server_info()
1148                    .expect("System status server should be started by DRT");
1149                let addr = system_info.socket_addr;
1150                let client = reqwest::Client::new();
1151                for (path, expect_200, expect_body) in [
1152                    ("/health", true, "ready"),
1153                    ("/live", true, "ready"),
1154                    ("/someRandomPathNotFoundHere", false, "Route not found"),
1155                ] {
1156                    println!("[test] Sending request to {}", path);
1157                    let url = format!("http://{}{}", addr, path);
1158                    let response = client.get(&url).send().await.unwrap();
1159                    let status = response.status();
1160                    let body = response.text().await.unwrap();
1161                    println!(
1162                        "[test] Response for {}: status={}, body={:?}",
1163                        path, status, body
1164                    );
1165                    if expect_200 {
1166                        assert_eq!(status, 200, "Response: status={}, body={:?}", status, body);
1167                    } else {
1168                        assert_eq!(status, 404, "Response: status={}, body={:?}", status, body);
1169                    }
1170                    assert!(
1171                        body.contains(expect_body),
1172                        "Response: status={}, body={:?}",
1173                        status,
1174                        body
1175                    );
1176                }
1177                // DRT handles server cleanup automatically
1178            },
1179        )
1180        .await;
1181    }
1182
1183    #[tokio::test]
1184    async fn test_payloadless_endpoint_is_healthy_with_canary_enabled() {
1185        temp_env::async_with_vars(
1186            [
1187                (env_system::DYN_SYSTEM_PORT, Some("0")),
1188                (
1189                    env_system::DYN_SYSTEM_STARTING_HEALTH_STATUS,
1190                    Some("notready"),
1191                ),
1192                ("DYN_HEALTH_CHECK_ENABLED", Some("true")),
1193            ],
1194            async {
1195                let runtime = crate::Runtime::from_current().unwrap();
1196                let drt = Arc::new(
1197                    crate::DistributedRuntime::new(
1198                        runtime,
1199                        crate::distributed::DistributedConfig::process_local(),
1200                    )
1201                    .await
1202                    .unwrap(),
1203                );
1204                let addr = drt
1205                    .system_status_server_info()
1206                    .expect("System status server should be started")
1207                    .socket_addr;
1208
1209                use crate::pipeline::{
1210                    AsyncEngine, AsyncEngineContextProvider, Error, ManyOut, SingleIn, async_trait,
1211                    network::Ingress,
1212                };
1213                use crate::protocols::annotated::Annotated;
1214
1215                struct PayloadlessHandler;
1216
1217                #[async_trait]
1218                impl AsyncEngine<SingleIn<String>, ManyOut<Annotated<String>>, anyhow::Error>
1219                    for PayloadlessHandler
1220                {
1221                    async fn generate(
1222                        &self,
1223                        input: SingleIn<String>,
1224                    ) -> anyhow::Result<ManyOut<Annotated<String>>> {
1225                        let (data, ctx) = input.into_parts();
1226                        Ok(crate::pipeline::ResponseStream::new(
1227                            Box::pin(crate::stream::iter(vec![Annotated::from_data(data)])),
1228                            ctx.context(),
1229                        ))
1230                    }
1231                }
1232
1233                let namespace = drt.namespace("test").unwrap();
1234                let component = namespace.component("backend").unwrap();
1235                let ingress = Ingress::for_engine(Arc::new(PayloadlessHandler)).unwrap();
1236                tokio::spawn(async move {
1237                    // Unified decode workers deliberately follow this path: the
1238                    // endpoint is registered without a canary payload.
1239                    let _ = component
1240                        .endpoint("generate")
1241                        .endpoint_builder()
1242                        .handler(ingress)
1243                        .start()
1244                        .await;
1245                });
1246
1247                let client = reqwest::Client::new();
1248                for path in ["/health", "/live"] {
1249                    let url = format!("http://{addr}{path}");
1250                    let mut last_response = None;
1251                    for _ in 0..50 {
1252                        let response = client.get(&url).send().await.unwrap();
1253                        let status = response.status();
1254                        let body = response.text().await.unwrap();
1255                        if status == 200 && body.contains("\"status\":\"ready\"") {
1256                            last_response = Some((status, body));
1257                            break;
1258                        }
1259                        last_response = Some((status, body));
1260                        tokio::time::sleep(Duration::from_millis(20)).await;
1261                    }
1262                    let (status, body) = last_response.unwrap();
1263                    assert_eq!(status, 200, "{path} response: {body}");
1264                    assert!(
1265                        body.contains("\"status\":\"ready\""),
1266                        "{path} response: {body}"
1267                    );
1268                }
1269            },
1270        )
1271        .await;
1272    }
1273
1274    #[cfg(feature = "integration")]
1275    #[tokio::test]
1276    async fn test_health_check_with_payload_and_timeout() {
1277        // Test the complete health check flow with the new canary-based system:
1278        crate::logging::init();
1279
1280        temp_env::async_with_vars(
1281            [
1282                (env_system::DYN_SYSTEM_PORT, Some("0")),
1283                (
1284                    env_system::DYN_SYSTEM_STARTING_HEALTH_STATUS,
1285                    Some("notready"),
1286                ),
1287                (
1288                    env_system::DYN_SYSTEM_USE_ENDPOINT_HEALTH_STATUS,
1289                    Some("[\"test.endpoint\"]"),
1290                ),
1291                // Enable health check with short intervals for testing
1292                ("DYN_HEALTH_CHECK_ENABLED", Some("true")),
1293                (env_canary::DYN_CANARY_WAIT_TIME, Some("1")), // Send canary after 1 second of inactivity
1294                ("DYN_HEALTH_CHECK_REQUEST_TIMEOUT", Some("1")), // Immediately timeout to mimic unresponsiveness
1295                ("RUST_LOG", Some("info")),                      // Enable logging for test
1296            ],
1297            async {
1298                let drt = Arc::new(create_test_drt_async().await);
1299
1300                // Get system status server info
1301                let system_info = drt
1302                    .system_status_server_info()
1303                    .expect("System status server should be started");
1304                let addr = system_info.socket_addr;
1305
1306                let client = reqwest::Client::new();
1307                let health_url = format!("http://{}/health", addr);
1308
1309                // Register an endpoint with health check payload
1310                let endpoint = "test.endpoint";
1311                let health_check_payload = serde_json::json!({
1312                    "prompt": "health check test",
1313                    "_health_check": true
1314                });
1315
1316                // Register the endpoint and its health check payload
1317                {
1318                    let system_health = drt.system_health();
1319                    let system_health_lock = system_health.lock();
1320                    system_health_lock.register_health_check_target(
1321                        endpoint,
1322                        crate::component::Instance {
1323                            component: "test_component".to_string(),
1324                            endpoint: "health".to_string(),
1325                            namespace: "test_namespace".to_string(),
1326                            instance_id: 1,
1327                            transport: crate::component::TransportType::Nats(endpoint.to_string()),
1328                            device_type: None,
1329                            request_plane_codec: None,
1330                        },
1331                        health_check_payload.clone(),
1332                    );
1333                }
1334
1335                // Check initial health - should be ready (default state)
1336                let response = client.get(&health_url).send().await.unwrap();
1337                let status = response.status();
1338                let body = response.text().await.unwrap();
1339                assert_eq!(status, 503, "Should be unhealthy initially (default state)");
1340                assert!(
1341                    body.contains("\"status\":\"notready\""),
1342                    "Should show notready status initially"
1343                );
1344
1345                // Set endpoint to healthy state
1346                drt.system_health()
1347                    .lock()
1348                    .set_endpoint_health_status(endpoint, HealthStatus::Ready);
1349
1350                // Check health again - should now be healthy
1351                let response = client.get(&health_url).send().await.unwrap();
1352                let status = response.status();
1353                let body = response.text().await.unwrap();
1354
1355                assert_eq!(status, 200, "Should be healthy due to recent response");
1356                assert!(
1357                    body.contains("\"status\":\"ready\""),
1358                    "Should show ready status after response"
1359                );
1360
1361                // Verify the endpoint status in SystemHealth directly
1362                let endpoint_status = drt
1363                    .system_health()
1364                    .lock()
1365                    .get_endpoint_health_status(endpoint);
1366                assert_eq!(
1367                    endpoint_status,
1368                    Some(HealthStatus::Ready),
1369                    "SystemHealth should show endpoint as Ready after response"
1370                );
1371            },
1372        )
1373        .await;
1374    }
1375
1376    /// `/v1/loras` compat shim: with the legacy local registry empty, a LoRA
1377    /// update registered in `engine_routes()` under `update/<name>` resolves via
1378    /// the fallback and its JSON response is parsed into a `LoraResponse`. This
1379    /// is the path unified-backend workers take for `/v1/loras`.
1380    #[tokio::test]
1381    async fn test_call_lora_endpoint_resolves_via_engine_routes() {
1382        temp_env::async_with_vars([(env_system::DYN_SYSTEM_PORT, None::<&str>)], async {
1383            let drt = create_test_drt_async().await;
1384
1385            let callback: crate::engine_routes::EngineRouteCallback = Arc::new(|_body| {
1386                Box::pin(async move {
1387                    Ok(serde_json::json!({
1388                        "status": "success",
1389                        "lora_name": "adapterA",
1390                        "lora_id": 42,
1391                    }))
1392                })
1393            });
1394            // Unified Worker registers LoRA ops under the `update/` namespace.
1395            drt.engine_routes().register("update/load_lora", callback);
1396
1397            // Local registry is empty, so resolution must fall through to
1398            // engine_routes.
1399            assert!(drt.local_endpoint_registry().get("load_lora").is_none());
1400
1401            let response = call_lora_endpoint(
1402                &drt,
1403                "load_lora",
1404                serde_json::json!({"lora_name": "adapterA"}),
1405            )
1406            .await
1407            .expect("engine_routes fallback should resolve the control");
1408
1409            assert_eq!(response.status, "success");
1410            assert_eq!(response.lora_name.as_deref(), Some("adapterA"));
1411            assert_eq!(response.lora_id, Some(42));
1412        })
1413        .await;
1414    }
1415
1416    /// When neither the local registry nor `engine_routes()` holds the name,
1417    /// the caller gets an explicit "LoRA management not available" error
1418    /// rather than an opaque "endpoint not found".
1419    #[tokio::test]
1420    async fn test_call_lora_endpoint_missing_returns_clean_error() {
1421        temp_env::async_with_vars([(env_system::DYN_SYSTEM_PORT, None::<&str>)], async {
1422            let drt = create_test_drt_async().await;
1423
1424            let err = call_lora_endpoint(&drt, "load_lora", serde_json::json!({}))
1425                .await
1426                .expect_err("missing handler must error");
1427
1428            assert!(
1429                err.to_string().contains("LoRA management not available"),
1430                "expected explicit unavailable message, got: {err}"
1431            );
1432        })
1433        .await;
1434    }
1435
1436    /// An update callback that returns `{"status":"error",...}` (rather than
1437    /// raising) surfaces as a `LoraResponse{status:"error"}`. The `/v1/loras`
1438    /// load/unload handlers map this to HTTP 500, preserving legacy semantics;
1439    /// the direct `/engine/*` route would instead return HTTP 200 + this JSON.
1440    #[tokio::test]
1441    async fn test_call_lora_endpoint_propagates_error_status() {
1442        temp_env::async_with_vars([(env_system::DYN_SYSTEM_PORT, None::<&str>)], async {
1443            let drt = create_test_drt_async().await;
1444
1445            let callback: crate::engine_routes::EngineRouteCallback = Arc::new(|_body| {
1446                Box::pin(async move {
1447                    Ok(serde_json::json!({
1448                        "status": "error",
1449                        "message": "adapter not found",
1450                    }))
1451                })
1452            });
1453            // Unified Worker registers LoRA ops under the `update/` namespace.
1454            drt.engine_routes().register("update/unload_lora", callback);
1455
1456            let response = call_lora_endpoint(&drt, "unload_lora", serde_json::json!({}))
1457                .await
1458                .expect("a non-raising callback returns Ok even on logical error");
1459
1460            assert_eq!(response.status, "error");
1461            assert_eq!(response.message.as_deref(), Some("adapter not found"));
1462        })
1463        .await;
1464    }
1465}
dynamo_runtime/system_status_server.rs

dynamo_runtime/
system_status_server.rs
