ForgeLM/crates/inference-server/src/routes.rs

use axum::{
    extract::{Json, State},
    http::{HeaderMap, StatusCode},
    response::{IntoResponse, Response, Sse},
    routing::{get, post},
    Router,
};
use std::{sync::Arc, time::SystemTime};
use tokio::sync::Mutex;
use tower::ServiceBuilder;
use tower_http::trace::{DefaultOnFailure, TraceLayer};

use api_contracts::{
    chat::{
        ChatCompletionChunk, ChatCompletionRequest, ChatCompletionResponse, ChatMessage, Choice,
        ChunkChoice, Delta, FinishReason, Role, Usage,
    },
    error::{ApiError, ErrorResponse},
    health::HealthResponse,
    models::ModelList,
};
use inference_core::InferenceEngine;
use metrics::counters::PerfCounters;

use crate::compat::{
    authorize_request, build_sampler, ensure_supported_mode, finalize_assistant_message,
    prepare_prompt, tool_calls_to_delta, validate_request, FinalizedAssistant, ServerConfig,
    StreamingMode,
};

#[derive(Clone)]
pub struct AppState {
    pub engine: Arc<Mutex<InferenceEngine>>,
    pub counters: Arc<PerfCounters>,
    pub model_name: Arc<String>,
    pub server_config: Arc<ServerConfig>,
    pub start_time: std::time::Instant,
}

pub fn build_router(state: AppState) -> Router {
    let cors = state.server_config.cors_layer();

    let trace_layer = TraceLayer::new_for_http()
        .make_span_with(|request: &axum::http::Request<_>| {
            let request_id = uuid::Uuid::new_v4().to_string();
            tracing::info_span!(
                "http",
                request_id = %request_id,
                method   = %request.method(),
                path     = %request.uri().path(),
                user_agent = request
                    .headers()
                    .get(axum::http::header::USER_AGENT)
                    .and_then(|v| v.to_str().ok())
                    .unwrap_or(""),
                // filled in on_response
                status   = tracing::field::Empty,
                latency_ms = tracing::field::Empty,
            )
        })
        .on_request(|request: &axum::http::Request<_>, _span: &tracing::Span| {
            let content_length = request
                .headers()
                .get(axum::http::header::CONTENT_LENGTH)
                .and_then(|v| v.to_str().ok())
                .unwrap_or("0");
            tracing::info!(content_length, "request_received");
        })
        .on_response(
            |response: &axum::http::Response<_>,
             latency: std::time::Duration,
             span: &tracing::Span| {
                span.record("status", response.status().as_u16());
                span.record("latency_ms", latency.as_millis());
                tracing::info!("response_sent");
            },
        )
        .on_failure(DefaultOnFailure::new().level(tracing::Level::ERROR));

    Router::new()
        .route("/health", get(health_handler))
        .route("/v1/models", get(models_handler))
        .route("/v1/chat/completions", post(chat_completions_handler))
        .layer(ServiceBuilder::new().layer(trace_layer).layer(cors))
        .with_state(state)
}

fn unix_now() -> i64 {
    SystemTime::now()
        .duration_since(SystemTime::UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs() as i64
}

fn error_response(status: StatusCode, error: ApiError) -> Response {
    (status, Json(ErrorResponse { error })).into_response()
}

/// Run `f` inside a `spawn_blocking` task while propagating the current tracing span.
/// Without this, logs emitted inside `spawn_blocking` lose the `request_id` context.
fn spawn_blocking_in_span<F, R>(f: F) -> tokio::task::JoinHandle<R>
where
    F: FnOnce() -> R + Send + 'static,
    R: Send + 'static,
{
    let span = tracing::Span::current();
    tokio::task::spawn_blocking(move || {
        let _guard = span.enter();
        f()
    })
}

async fn health_handler(State(state): State<AppState>) -> impl IntoResponse {
    Json(HealthResponse {
        status: "ok",
        model_loaded: true,
        model_id: Some(state.model_name.as_ref().clone()),
        uptime_secs: state.start_time.elapsed().as_secs(),
    })
}

async fn models_handler(State(state): State<AppState>, headers: HeaderMap) -> Response {
    if let Err((status, error)) = authorize_request(&headers, &state.server_config) {
        tracing::warn!(status = %status, "models: authorization failed");
        return error_response(status, error);
    }

    let model_list = state.server_config.model_list(unix_now());
    tracing::info!(model_count = model_list.len(), "models.listed");
    Json(ModelList { object: "list", data: model_list }).into_response()
}

async fn chat_completions_handler(
    State(state): State<AppState>,
    headers: HeaderMap,
    Json(req): Json<ChatCompletionRequest>,
) -> Response {
    tracing::info!(
        model       = %req.model,
        stream      = req.stream,
        messages    = req.messages.len(),
        max_tokens  = ?req.max_tokens,
        temperature = ?req.temperature,
        auth_present = headers.contains_key(
            state.server_config.auth_header
                .parse::<axum::http::HeaderName>()
                .unwrap_or(axum::http::header::AUTHORIZATION)
        ),
        "chat_completions.received"
    );

    if let Err((status, error)) = authorize_request(&headers, &state.server_config) {
        tracing::warn!(status = %status, "chat_completions: authorization failed");
        return error_response(status, error);
    }

    if let Err((status, error)) = validate_request(&req, &state.server_config) {
        tracing::warn!(status = %status, "chat_completions: validation failed");
        return error_response(status, error);
    }

    if let Err(err) = ensure_supported_mode(&req) {
        tracing::warn!(err = %err, "chat_completions: unsupported mode");
        return error_response(
            StatusCode::BAD_REQUEST,
            ApiError::new(err.to_string(), "invalid_request_error").with_code("unsupported_mode"),
        );
    }

    let sampler = build_sampler(&req);
    let prepared = prepare_prompt(&req, &state.server_config);
    tracing::debug!(prompt_len = prepared.prompt.len(), "chat_completions: prompt prepared");

    if req.stream {
        tracing::info!("chat_completions: starting stream response");
        stream_response(
            state,
            req,
            sampler,
            prepared.prompt,
            prepared.streaming_mode,
        )
        .await
    } else {
        tracing::info!("chat_completions: starting blocking response");
        blocking_response(state, req, sampler, prepared.prompt).await
    }
}

async fn stream_response(
    state: AppState,
    req: ChatCompletionRequest,
    sampler: inference_core::SamplerConfig,
    prompt: String,
    streaming_mode: StreamingMode,
) -> Response {
    use axum::response::sse::{Event, KeepAlive};

    let id = uuid::Uuid::new_v4().to_string();
    let created = unix_now();
    let resolved_model = state
        .server_config
        .resolve_model(&req.model)
        .map(|model| model.root.clone())
        .unwrap_or_else(|| state.model_name.as_ref().clone());

    let engine_arc = Arc::clone(&state.engine);
    let counters_arc = Arc::clone(&state.counters);

    let stream = async_stream::stream! {
        let first_chunk = ChatCompletionChunk {
            id: id.clone(),
            object: "chat.completion.chunk",
            created,
            model: resolved_model.clone(),
            choices: vec![ChunkChoice {
                index: 0,
                delta: Delta { role: Some(Role::Assistant), content: None, tool_calls: None },
                finish_reason: None,
            }],
            usage: None,
        };
        yield Ok::<_, std::convert::Infallible>(Event::default().data(serde_json::to_string(&first_chunk).unwrap()));

        match streaming_mode {
            StreamingMode::LiveTokens => {
                let (tx, mut rx) = tokio::sync::mpsc::channel::<StreamItem>(256);
                let prompt = prompt.clone();

                spawn_blocking_in_span(move || {
                    let rt = tokio::runtime::Handle::current();
                    tracing::debug!("stream: acquiring engine lock");
                    let mut engine = rt.block_on(engine_arc.lock());
                    tracing::debug!("stream: engine lock acquired, creating session");
                    let mut record = metrics::recorder::GenerationRecord::new(0);

                    let mut session = match engine.new_session(&prompt, sampler) {
                        Ok(session) => {
                            tracing::info!(prompt_tokens = session.prompt_tokens(), "stream: session created");
                            session
                        }
                        Err(err) => {
                            tracing::error!(err = %err, "Failed to create session");
                            let _ = rt.block_on(tx.send(StreamItem::Error));
                            return;
                        }
                    };

                    while let Some(token) = session.next_token() {
                        record.tick();
                        counters_arc.inc_tokens(1);
                        let _ = rt.block_on(tx.send(StreamItem::Token(token)));
                    }

                    let usage = Usage {
                        prompt_tokens: session.prompt_tokens(),
                        completion_tokens: record.completion_tokens,
                        total_tokens: session.prompt_tokens() + record.completion_tokens,
                        tokens_per_second: record.tokens_per_second(),
                    };
                    tracing::info!(
                        prompt_tokens     = usage.prompt_tokens,
                        completion_tokens = usage.completion_tokens,
                        total_tokens      = usage.total_tokens,
                        tps               = ?usage.tokens_per_second,
                        finish_reason     = "stop",
                        "stream.generation_complete"
                    );
                    counters_arc.inc_completed();
                    let _ = rt.block_on(tx.send(StreamItem::Done {
                        finish_reason: FinishReason::Stop,
                        usage,
                    }));
                });

                let mut streamed_tokens: u32 = 0;
                while let Some(item) = rx.recv().await {
                    match item {
                        StreamItem::Token(token) => {
                            streamed_tokens += 1;
                            if streamed_tokens <= 3 || streamed_tokens % 50 == 0 {
                                tracing::debug!(n = streamed_tokens, token = %token, "stream: token");
                            }
                            let chunk = ChatCompletionChunk {
                                id: id.clone(),
                                object: "chat.completion.chunk",
                                created,
                                model: resolved_model.clone(),
                                choices: vec![ChunkChoice {
                                    index: 0,
                                    delta: Delta { role: None, content: Some(token), tool_calls: None },
                                    finish_reason: None,
                                }],
                                usage: None,
                            };
                            yield Ok(Event::default().data(serde_json::to_string(&chunk).unwrap()));
                        }
                        StreamItem::Done { finish_reason, usage } => {
                            // prompt_tokens/completion_tokens/finish_reason already logged
                            // in stream.generation_complete inside spawn_blocking.
                            tracing::info!(streamed_tokens, "stream.final_chunk");
                            let final_chunk = ChatCompletionChunk {
                                id: id.clone(),
                                object: "chat.completion.chunk",
                                created,
                                model: resolved_model.clone(),
                                choices: vec![ChunkChoice {
                                    index: 0,
                                    delta: Delta { role: None, content: None, tool_calls: None },
                                    finish_reason: Some(finish_reason),
                                }],
                                usage: req.stream_options.as_ref().filter(|options| options.include_usage).map(|_| usage),
                            };
                            yield Ok(Event::default().data(serde_json::to_string(&final_chunk).unwrap()));
                            yield Ok(Event::default().data("[DONE]"));
                            break;
                        }
                        StreamItem::Error => {
                            tracing::error!(streamed_tokens, "stream.error: session failed — sending empty finish chunk");
                            let chunk = ChatCompletionChunk {
                                id: id.clone(),
                                object: "chat.completion.chunk",
                                created,
                                model: resolved_model.clone(),
                                choices: vec![ChunkChoice {
                                    index: 0,
                                    delta: Delta { role: None, content: None, tool_calls: None },
                                    finish_reason: Some(FinishReason::Stop),
                                }],
                                usage: None,
                            };
                            yield Ok(Event::default().data(serde_json::to_string(&chunk).unwrap()));
                            yield Ok(Event::default().data("[DONE]"));
                            break;
                        }
                    }
                }
            }
            StreamingMode::Buffered => {
                match generate_full_completion(engine_arc, counters_arc, prompt, sampler, req.clone()).await {
                    Ok(result) => {
                        if let Some(tool_calls) = result.assistant.tool_calls {
                            let chunk = ChatCompletionChunk {
                                id: id.clone(),
                                object: "chat.completion.chunk",
                                created,
                                model: resolved_model.clone(),
                                choices: vec![ChunkChoice {
                                    index: 0,
                                    delta: Delta {
                                        role: None,
                                        content: None,
                                        tool_calls: Some(tool_calls_to_delta(&tool_calls)),
                                    },
                                    finish_reason: None,
                                }],
                                usage: None,
                            };
                            yield Ok(Event::default().data(serde_json::to_string(&chunk).unwrap()));
                        } else if let Some(content) = result.assistant.content {
                            let chunk = ChatCompletionChunk {
                                id: id.clone(),
                                object: "chat.completion.chunk",
                                created,
                                model: resolved_model.clone(),
                                choices: vec![ChunkChoice {
                                    index: 0,
                                    delta: Delta { role: None, content: Some(content), tool_calls: None },
                                    finish_reason: None,
                                }],
                                usage: None,
                            };
                            yield Ok(Event::default().data(serde_json::to_string(&chunk).unwrap()));
                        }

                        let final_chunk = ChatCompletionChunk {
                            id: id.clone(),
                            object: "chat.completion.chunk",
                            created,
                            model: resolved_model.clone(),
                            choices: vec![ChunkChoice {
                                index: 0,
                                delta: Delta { role: None, content: None, tool_calls: None },
                                finish_reason: Some(result.assistant.finish_reason),
                            }],
                            usage: req.stream_options.as_ref().filter(|options| options.include_usage).map(|_| result.usage),
                        };
                        yield Ok(Event::default().data(serde_json::to_string(&final_chunk).unwrap()));
                        yield Ok(Event::default().data("[DONE]"));
                    }
                    Err(_response) => {
                        tracing::error!("stream.buffered: generation failed");
                        yield Ok(Event::default().data("[DONE]"));
                    }
                }
            }
        }
    };

    let response = Sse::new(stream)
        .keep_alive(KeepAlive::default())
        .into_response();

    axum::response::Response::builder()
        .status(response.status())
        .header("content-type", "text/event-stream; charset=utf-8")
        .header("cache-control", "no-cache")
        .header("connection", "keep-alive")
        .body(response.into_body())
        .unwrap()
}

async fn blocking_response(
    state: AppState,
    req: ChatCompletionRequest,
    sampler: inference_core::SamplerConfig,
    prompt: String,
) -> Response {
    let id = uuid::Uuid::new_v4().to_string();
    let model = state
        .server_config
        .resolve_model(&req.model)
        .map(|descriptor| descriptor.root.clone())
        .unwrap_or_else(|| state.model_name.as_ref().clone());

    match generate_full_completion(
        Arc::clone(&state.engine),
        Arc::clone(&state.counters),
        prompt,
        sampler,
        req,
    )
    .await
    {
        Ok(result) => Json(ChatCompletionResponse {
            id,
            object: "chat.completion",
            created: unix_now(),
            model,
            choices: vec![Choice {
                index: 0,
                message: ChatMessage {
                    role: Role::Assistant,
                    content: result.assistant.content,
                    name: None,
                    tool_calls: result.assistant.tool_calls,
                    tool_call_id: None,
                },
                finish_reason: result.assistant.finish_reason,
            }],
            usage: result.usage,
        })
        .into_response(),
        Err(response) => response,
    }
}

async fn generate_full_completion(
    engine_arc: Arc<Mutex<InferenceEngine>>,
    counters_arc: Arc<PerfCounters>,
    prompt: String,
    sampler: inference_core::SamplerConfig,
    req: ChatCompletionRequest,
) -> Result<GenerationResult, Response> {
    tracing::debug!(prompt_len = prompt.len(), "generate_full_completion: starting");
    let result = spawn_blocking_in_span(move || {
        let rt = tokio::runtime::Handle::current();
        tracing::debug!("generate_full_completion: acquiring engine lock");
        let mut engine = rt.block_on(engine_arc.lock());
        tracing::debug!("generate_full_completion: engine lock acquired");
        let mut record = metrics::recorder::GenerationRecord::new(0);
        let mut content = String::new();

        let mut session = engine.new_session(&prompt, sampler)?;
        tracing::info!(prompt_tokens = session.prompt_tokens(), "generate_full_completion: session created");

        while let Some(tok) = session.next_token() {
            record.tick();
            content.push_str(&tok);
        }

        counters_arc.inc_tokens(record.completion_tokens as u64);
        counters_arc.inc_completed();

        let usage = Usage {
            prompt_tokens: session.prompt_tokens(),
            completion_tokens: record.completion_tokens,
            total_tokens: session.prompt_tokens() + record.completion_tokens,
            tokens_per_second: record.tokens_per_second(),
        };
        tracing::info!(
            prompt_tokens = usage.prompt_tokens,
            completion_tokens = usage.completion_tokens,
            content_len = content.len(),
            tps = ?usage.tokens_per_second,
            "generate_full_completion: done"
        );
        let assistant = finalize_assistant_message(content, &req);
        Ok::<_, anyhow::Error>((assistant, usage))
    })
    .await;

    match result {
        Ok(Ok((assistant, usage))) => Ok(GenerationResult { assistant, usage }),
        Ok(Err(err)) => {
            tracing::error!(err = %err, "Generation failed");
            Err(error_response(
                StatusCode::INTERNAL_SERVER_ERROR,
                ApiError::new("Generation failed", "server_error").with_code("generation_failed"),
            ))
        }
        Err(err) => {
            tracing::error!(err = %err, "Spawn blocking panicked");
            Err(error_response(
                StatusCode::INTERNAL_SERVER_ERROR,
                ApiError::new("Internal server error", "server_error").with_code("internal_error"),
            ))
        }
    }
}

struct GenerationResult {
    assistant: FinalizedAssistant,
    usage: Usage,
}

enum StreamItem {
    Token(String),
    Done {
        finish_reason: FinishReason,
        usage: Usage,
    },
    Error,
}