// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0

//! Tests that the router does not wedge indefinitely when an upstream
//! worker accepts the TCP connection but never sends response headers.
//!
//! Without a configured `.timeout(...)` on the reqwest client, a stalled
//! backend hangs the axum handler future forever and the test harness
//! would just timeout. We assert here that the router returns a fast,
//! clean 504 (`upstream_timeout`) instead — a timeout is a gateway timeout,
//! the same status class as the stale-deadline cancel.

use axum::body::Body;
use axum::http::{Request, StatusCode};
use http_body_util::BodyExt;
use sgl_router::config::{
    Config, DiscoveryBackend, InflightLoadConfig, ModelConfig, ObservabilityConfig, PolicyKind,
    ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
};
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
use sgl_router::policies::factory::build_registry_with_defaults as build_policy_registry;
use sgl_router::proxy::Proxy;
use sgl_router::server::app::build_router;
use sgl_router::server::app_context::AppContext;
use sgl_router::tokenizer::TokenizerRegistry;
use sgl_router::workers::WorkerRegistry;
use std::sync::Arc;
use std::time::Duration;
use tower::ServiceExt;

fn config(_worker_url: &str) -> Config {
    Config {
        server: ServerConfig {
            host: "0".into(),
            port: 0,
            ..Default::default()
        },
        observability: ObservabilityConfig::default(),
        model: ModelConfig {
            id: "tiny".into(),
            tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
            disable_input_ids_forwarding: false,
            policy: PolicyKind::RoundRobin,
            decode_policy: Default::default(),
            bucket_config: None,
            circuit_breaker: None,
            cache_aware: None,
            sticky: None,
            affinity: None,
            fused: None,
            eligibility: None,
            sampling_overrides: Default::default(),
        },
        discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
            urls: vec!["http://placeholder:0".into()],
        }),
        proxy: ProxyConfig::default(),
        router_inflight_load: InflightLoadConfig::default(),
    }
}

#[tokio::test]
async fn non_streaming_request_times_out_when_worker_hangs() {
    // Worker accepts and then sleeps for 5s; router timeout is 200ms.
    let worker =
        crate::common::mock_worker::MockWorker::start_hanging(Duration::from_secs(5)).await;
    let cfg = config(&worker.url);
    let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
    let registry = Arc::new(WorkerRegistry::default());
    let _ = registry.add(WorkerSpec {
        id: WorkerId("w1".into()),
        url: worker.url.clone(),
        mode: WorkerMode::Plain,
        model_ids: vec![ModelId("tiny".into())],
        bootstrap_port: None,
    });
    let policies = Arc::new(build_policy_registry(&cfg).unwrap());
    let proxy = Arc::new(Proxy::new(Duration::from_millis(200)).unwrap());
    let ctx = Arc::new(AppContext::new(cfg, tokenizers, proxy, registry, policies));
    let app = build_router(ctx.clone());

    let req = Request::builder()
        .method("POST")
        .uri("/v1/chat/completions")
        .header("content-type", "application/json")
        .body(Body::from(
            serde_json::to_vec(&serde_json::json!({
                "model": "tiny",
                "messages": [{"role": "user", "content": "hi"}],
                "stream": false
            }))
            .unwrap(),
        ))
        .unwrap();

    let started = std::time::Instant::now();
    // Outer guard so a regression doesn't wedge CI forever.
    let res = tokio::time::timeout(Duration::from_secs(2), app.oneshot(req))
        .await
        .expect("router must return within 2s when proxy timeout is 200ms")
        .unwrap();
    let elapsed = started.elapsed();
    assert!(
        elapsed < Duration::from_secs(1),
        "router must short-circuit on upstream timeout; elapsed {elapsed:?}"
    );
    assert_eq!(res.status(), StatusCode::GATEWAY_TIMEOUT);
    assert_eq!(
        res.headers().get("x-router-error-code").unwrap(),
        "upstream_timeout"
    );
    let bytes = res.into_body().collect().await.unwrap().to_bytes();
    let body_str = String::from_utf8_lossy(&bytes);
    // A hung worker is the most common hard worker failure there is, so it must
    // land in `outcome="error"` — the series a per-worker error-ratio alert
    // watches. Deriving the outcome from the 504 status instead would silently
    // reclassify it as `cancelled` and blind that alert.
    assert!(
        ctx.metrics
            .render()
            .lines()
            .any(|l| l.starts_with("sgl_router_worker_requests_total{")
                && l.contains(r#"outcome="error""#)),
        "an upstream timeout must be counted outcome=error, not cancelled:\n{}",
        ctx.metrics.render(),
    );
    assert!(
        body_str.contains("\"code\":\"upstream_timeout\""),
        "body: {body_str}"
    );
    // No leak of worker URL or reqwest source chain to the client.
    assert!(
        !body_str.contains(&worker.url),
        "worker URL must not leak in client-visible body: {body_str}"
    );
}