rustytorch/crates/production/rtx-serving-api/src/server.rs

//! HTTP server configuration and startup

use axum::{
    Router,
    routing::{get, post},
};
use std::time::Duration;
use tokio::net::TcpListener;
use tower_http::{cors::CorsLayer, timeout::TimeoutLayer, trace::TraceLayer};

use crate::{ApiError, ApiResult, health, inference, models};

/// Server configuration
#[derive(Debug, Clone)]
pub struct ServerConfig {
    /// Listen address
    pub host: String,
    /// Listen port
    pub port: u16,
    /// Request timeout
    pub timeout_seconds: u64,
}

impl Default for ServerConfig {
    fn default() -> Self {
        Self {
            host: "127.0.0.1".to_string(),
            port: 8080,
            timeout_seconds: 30,
        }
    }
}

/// Serving server
pub struct ServingServer {
    config: ServerConfig,
    state: inference::AppState,
}

impl ServingServer {
    /// Create a new serving server (no engine loaded; inference returns 503)
    #[must_use]
    pub fn new(config: ServerConfig) -> Self {
        Self {
            config,
            state: inference::AppState::default(),
        }
    }

    /// Create a serving server backed by a live inference engine.
    ///
    /// Uses byte-level tokenization by default; use
    /// [`ServingServer::with_engine_and_tokenizer`] to attach a real
    /// vocabulary tokenizer.
    #[must_use]
    pub fn with_engine(
        config: ServerConfig,
        engine: std::sync::Arc<tokio::sync::RwLock<rtx_inference::InferenceEngine>>,
    ) -> Self {
        Self::with_engine_and_tokenizer(config, engine, rtx_inference::ServingTokenizer::default())
    }

    /// Create a serving server backed by a live inference engine and an
    /// explicit tokenizer (e.g. loaded via
    /// `rtx_inference::ServingTokenizer::from_file`).
    #[must_use]
    pub fn with_engine_and_tokenizer(
        config: ServerConfig,
        engine: std::sync::Arc<tokio::sync::RwLock<rtx_inference::InferenceEngine>>,
        tokenizer: rtx_inference::ServingTokenizer,
    ) -> Self {
        Self {
            config,
            state: inference::AppState {
                engine: Some(engine),
                tokenizer,
            },
        }
    }

    /// Build the Axum router with all routes
    pub fn app(&self) -> Router {
        Router::new()
            // Health endpoints
            .route("/health", get(health::health_check))
            .route("/health/ready", get(health::health_check))
            .route("/health/live", get(health::health_check))
            // Model endpoints
            .route("/v1/models", get(models::list_models))
            // Inference endpoints
            .route("/v1/completions", post(inference::inference))
            .route("/v1/chat/completions", post(inference::inference))
            // Enhanced cached inference endpoints (temporarily disabled)
            // .route("/v1/cached/completions", post(inference_cached::cached_inference))
            // .route("/v1/cached/chat/completions", post(inference_cached::cached_inference))
            // Cache management endpoints (temporarily disabled)
            // .route("/v1/cache/stats", get(inference_cached::cache_stats))
            // .route("/v1/cache/clear", post(inference_cached::clear_cache))
            // .route("/v1/cache/warm", post(inference_cached::warm_cache))
            // Middleware
            .layer(CorsLayer::permissive())
            .layer(TraceLayer::new_for_http())
            .layer(TimeoutLayer::new(Duration::from_secs(
                self.config.timeout_seconds,
            )))
            .with_state(self.state.clone())
    }

    /// Start the server
    pub async fn serve(&self) -> ApiResult<()> {
        let app = self.app();

        let addr = format!("{}:{}", self.config.host, self.config.port);
        let listener = TcpListener::bind(&addr)
            .await
            .map_err(|e| ApiError::internal(format!("Failed to bind to {addr}: {e}")))?;

        tracing::info!("Server listening on {}", addr);

        axum::serve(listener, app)
            .await
            .map_err(|e| ApiError::internal(format!("Server error: {e}")))?;

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_server_config_default() {
        let config = ServerConfig::default();
        assert_eq!(config.host, "127.0.0.1");
        assert_eq!(config.port, 8080);
        assert_eq!(config.timeout_seconds, 30);
    }

    #[test]
    fn test_server_creation() {
        let config = ServerConfig::default();
        let server = ServingServer::new(config);

        // Should be able to build the app router
        let _app = server.app();
    }
}